diff --git a/src/biotite/sequence/align/buckets.py b/src/biotite/sequence/align/buckets.py index b8159a0cb..df766ba39 100644 --- a/src/biotite/sequence/align/buckets.py +++ b/src/biotite/sequence/align/buckets.py @@ -6,10 +6,13 @@ __author__ = "Patrick Kunzmann" __all__ = ["bucket_number"] -from os.path import dirname, join, realpath +import functools +from pathlib import Path +from typing import Any import numpy as np +from biotite.typing import NDArray1 -_primes = None +_PRIMES_PATH = Path(__file__).parent / "primes.txt" def bucket_number(n_kmers: int, load_factor: float = 0.8) -> int: @@ -53,19 +56,22 @@ def bucket_number(n_kmers: int, load_factor: float = 0.8) -> int: Hence, all *k-mers* with the same suffix would be stored in the same bin. """ - global _primes - if _primes is None: - with open(join(dirname(realpath(__file__)), "primes.txt")) as file: - _primes = np.array( - [ - int(line) - for line in file.read().splitlines() - if len(line) != 0 and line[0] != "#" - ] - ) - + primes = _get_primes() number = int(n_kmers / load_factor) - index = np.searchsorted(_primes, number, side="left") - if index == len(_primes): + index = np.searchsorted(primes, number, side="left") + if index == len(primes): raise ValueError("Number of buckets too large") - return _primes[index] + return primes[index] + + +@functools.cache +def _get_primes() -> NDArray1[Any, np.integer]: + with open(_PRIMES_PATH) as file: + _primes = np.array( + [ + int(line) + for line in file.read().splitlines() + if len(line) != 0 and line[0] != "#" + ] + ) + return _primes diff --git a/src/biotite/sequence/align/kmertable.py b/src/biotite/sequence/align/kmertable.py new file mode 100644 index 000000000..5ce7f55d2 --- /dev/null +++ b/src/biotite/sequence/align/kmertable.py @@ -0,0 +1,16 @@ +# 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__ = ["KmerTable", "BucketKmerTable"] + +from biotite.rust.sequence.align import BucketKmerTable, KmerTable + +# The classes are implemented in Rust and therefore are not generic at runtime; +# allow ``KmerTable[S]`` subscription so the annotated stub works +KmerTable.__class_getitem__ = classmethod(lambda cls, params: cls) +BucketKmerTable.__class_getitem__ = classmethod(lambda cls, params: cls) diff --git a/src/biotite/sequence/align/kmertable.pyi b/src/biotite/sequence/align/kmertable.pyi index cb3128462..0d23e572b 100644 --- a/src/biotite/sequence/align/kmertable.pyi +++ b/src/biotite/sequence/align/kmertable.pyi @@ -1,7 +1,7 @@ __all__ = ["KmerTable", "BucketKmerTable"] from collections.abc import Iterable, Iterator -from typing import Any, Generic +from typing import Generic import numpy as np from biotite.sequence.align.kmeralphabet import KmerAlphabet from biotite.sequence.align.kmersimilarity import SimilarityRule @@ -10,7 +10,6 @@ from biotite.sequence.sequence import Sequence from biotite.typing import C2, C3, C4, M, N, NDArray1, NDArray2, S class KmerTable(Generic[S]): - def __init__(self, kmer_alphabet: KmerAlphabet[S]) -> None: ... @property def kmer_alphabet(self) -> KmerAlphabet[S]: ... @property @@ -74,12 +73,8 @@ class KmerTable(Generic[S]): def __reversed__(self) -> Iterator[int]: ... def __eq__(self, item: object) -> bool: ... def __str__(self) -> str: ... - def __getnewargs_ex__(self) -> tuple[tuple, dict]: ... - def __getstate__(self) -> Any: ... - def __setstate__(self, state: Any) -> None: ... class BucketKmerTable(Generic[S]): - def __init__(self, n_buckets: int, kmer_alphabet: KmerAlphabet[S]) -> None: ... @property def kmer_alphabet(self) -> KmerAlphabet[S]: ... @property @@ -140,6 +135,3 @@ class BucketKmerTable(Generic[S]): def __len__(self) -> int: ... def __eq__(self, item: object) -> bool: ... def __str__(self) -> str: ... - def __getnewargs_ex__(self) -> tuple[tuple, dict]: ... - def __getstate__(self) -> Any: ... - def __setstate__(self, state: Any) -> None: ... diff --git a/src/biotite/sequence/align/kmertable.pyx b/src/biotite/sequence/align/kmertable.pyx deleted file mode 100644 index 67759bcd6..000000000 --- a/src/biotite/sequence/align/kmertable.pyx +++ /dev/null @@ -1,3411 +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. - -# distutils: language = c++ - -__name__ = "biotite.sequence.align" -__author__ = "Patrick Kunzmann" -__all__ = ["KmerTable", "BucketKmerTable"] - -cimport cython -cimport numpy as np -from cpython.mem cimport PyMem_Malloc as malloc -from cpython.mem cimport PyMem_Free as free -from libc.string cimport memcpy -from libcpp.set cimport set as cpp_set - -import numpy as np -from ..alphabet import LetterAlphabet, common_alphabet, AlphabetError -from .kmeralphabet import KmerAlphabet -from .buckets import bucket_number - - -ctypedef np.int32_t int32 -ctypedef np.int64_t int64 -ctypedef np.uint8_t uint8 -ctypedef np.uint32_t uint32 -ctypedef np.uint64_t ptr - - -cdef enum EntrySize: - # The size (number of 32 bit elements) for each entry in C-arrays - # of KmerTable and BucketKmerTable, respectively - # - # Size: reference ID (int32) + sequence pos (int32) - NO_BUCKETS = 2 - # Size: k-mer (int64) + reference ID (int32) + sequence pos (int32) - BUCKETS = 4 - - -cdef class KmerTable: - """ - This class represents a *k-mer* index table. - It maps *k-mers* (subsequences with length *k*) to the sequence - positions, where the *k-mer* appears. - It is primarily used to find *k-mer* matches between two sequences. - A match is defined as a *k-mer* that appears in both sequences. - Instances of this class are immutable. - - There are multiple ways to create a :class:`KmerTable`: - - - :meth:`from_sequences()` iterates through all overlapping - *k-mers* in a sequence and stores the sequence position of - each *kmer* in the table. - - :meth:`from_kmers()` is similar to :meth:`from_sequences()` - but directly accepts *k-mers* as input instead of sequences. - - :meth:`from_kmer_selection()` takes a combination of *k-mers* - and their positions in a sequence, which can be used to - apply subset selectors, such as :class:`MinimizerSelector`. - - :meth:`from_tables()` merges the entries from multiple - :class:`KmerTable` objects into a new table. - - :meth:`from_positions()` let's the user provide manual - *k-mer* positions, which can be useful for loading a - :class:`KmerTable` from file. - - The standard constructor merely returns an empty table and is - reserved for internal use. - - Each indexed *k-mer* position is represented by a tuple of - - 1. a unique reference ID that identifies to which sequence a - position refers to and - 2. the zero-based sequence position of the first symbol in the - *k-mer*. - - The :meth:`match()` method iterates through all overlapping *k-mers* - in another sequence and, for each *k-mer*, looks up the reference - IDs and positions of this *k-mer* in the table. - For each matching position, it adds the *k-mer* position in this - sequence, the matching reference ID and the matching sequence - position to the array of matches. - Finally these matches are returned to the user. - Optionally, a :class:`SimilarityRule` can be supplied, to find - also matches for similar *k-mers*. - This is especially useful for protein sequences to match two - *k-mers* with a high substitution probability. - - The positions for a given *k-mer* code can be obtained via indexing. - Iteration over a :class:`KmerTable` yields the *k-mers* that have at - least one associated position. - The *k-mer* code for a *k-mer* can be calculated with - ``table.kmer_alphabet.encode()`` (see :class:`KmerAlphabet`). - - Attributes - ---------- - kmer_alphabet : KmerAlphabet - The internal :class:`KmerAlphabet`, that is used to - encode all overlapping *k-mers* of an input sequence. - alphabet : Alphabet - The base alphabet, from which this :class:`KmerTable` was - created. - k : int - The length of the *k-mers*. - - See Also - -------- - BucketKmerTable - - Notes - ----- - - The design of the :class:`KmerTable` is inspired by the *MMseqs2* - software :footcite:`Steinegger2017`. - - *Memory consumption* - - For efficient mapping, a :class:`KmerTable` contains a pointer - array, that contains one 64-bit pointer for each possible *k-mer*. - If there is at least one position for a *k-mer*, the corresponding - pointer points to a C-array that contains - - 1. The length of the C-array *(int64)* - 2. The reference ID for each position of this *k-mer* *(uint32)* - 3. The sequence position for each position of this *k-mer* *(uint32)* - - Hence, the memory requirements can be quite large for long *k-mers* - or large alphabets. - The required memory space :math:`S` in byte is within the bounds of - - .. math:: - - 8 n^k + 8L \leq S \leq 16 n^k + 8L, - - where :math:`n` is the number of symbols in the alphabet and - :math:`L` is the summed length of all sequences added to the table. - - *Multiprocessing* - - :class:`KmerTable` objects can be used in multi-processed setups: - Adding a large database of sequences to a table can be sped up by - splitting the database into smaller chunks and create a separate - table for each chunk in separate processes. - Eventually, the tables can be merged to one large table using - :meth:`from_tables()`. - - Since :class:`KmerTable` supports the *pickle* protocol, - the matching step can also be divided into multiple processes, if - multiple sequences need to be matched. - - *Storage on hard drive* - - The most time efficient way to read/write a :class:`KmerTable` is - the *pickle* format. - If a custom format is desired, the user needs to extract the - reference IDs and position for each *k-mer*. - To restrict this task to all *k-mer* that have at least one match - :meth:`get_kmers()` can be used. - Conversely, the reference IDs and positions can be restored via - :meth:`from_positions()`. - - References - ---------- - - .. footbibliography:: - - Examples - -------- - - Create a *2-mer* index table for some nucleotide sequences: - - >>> table = KmerTable.from_sequences( - ... k = 2, - ... sequences = [NucleotideSequence("TTATA"), NucleotideSequence("CTAG")], - ... ref_ids = [0, 1] - ... ) - - Display the contents of the table as - (reference ID, sequence position) tuples: - - >>> print(table) - AG: (1, 2) - AT: (0, 2) - CT: (1, 0) - TA: (0, 1), (0, 3), (1, 1) - TT: (0, 0) - - Find matches of the table with a sequence: - - >>> query = NucleotideSequence("TAG") - >>> matches = table.match(query) - >>> for query_pos, table_ref_id, table_pos in matches: - ... print("Query sequence position:", query_pos) - ... print("Table reference ID: ", table_ref_id) - ... print("Table sequence position:", table_pos) - ... print() - Query sequence position: 0 - Table reference ID: 0 - Table sequence position: 1 - - Query sequence position: 0 - Table reference ID: 0 - Table sequence position: 3 - - Query sequence position: 0 - Table reference ID: 1 - Table sequence position: 1 - - Query sequence position: 1 - Table reference ID: 1 - Table sequence position: 2 - - - Get all reference IDs and positions for a given *k-mer*: - - >>> kmer_code = table.kmer_alphabet.encode("TA") - >>> print(table[kmer_code]) - [[0 1] - [0 3] - [1 1]] - """ - - cdef object _kmer_alph - cdef int _k - - # The pointer array is the core of the index table: - # It maps each possible k-mer (represented by its code) to a - # C-array of indices. - # Each entry in a C-array points to a reference ID and the - # location in that sequence where the respective k-mer appears - # The memory layout of each C-array is as following: - # - # (Array length) (RefID 0) (Position 0) (RefID 1) (Position 1) ... - # -----int64----|---uint32---|---uint32---|---uint32---|---uint32--- - # - # The array length is based on 32 bit units. - # If there is no entry for a k-mer, the respective pointer is NULL. - cdef ptr[:] _ptr_array - - - def __cinit__(self, kmer_alphabet): - # This check is necessary for proper memory management - # of the allocated arrays - if self._is_initialized(): - raise Exception("Duplicate call of constructor") - - self._kmer_alph = kmer_alphabet - self._k = kmer_alphabet.k - self._ptr_array = np.zeros(len(self._kmer_alph), dtype=np.uint64) - - - @property - def kmer_alphabet(self): - return self._kmer_alph - - @property - def alphabet(self): - return self._kmer_alph.base_alphabet - - @property - def k(self): - return self._k - - @staticmethod - def from_sequences(k, sequences, ref_ids=None, ignore_masks=None, - alphabet=None, spacing=None): - """ - from_sequences(k, sequences, ref_ids=None, ignore_masks=None, - alphabet=None, spacing=None) - - Create a :class:`KmerTable` by storing the positions of all - overlapping *k-mers* from the input `sequences`. - - Parameters - ---------- - k : int - The length of the *k-mers*. - sequences : sized iterable object of Sequence, length=m - The sequences to get the *k-mer* positions from. - These sequences must have equal alphabets, or one of these - sequences must have an alphabet that extends the alphabets - of all other sequences. - ref_ids : sized iterable object of int, length=m, optional - The reference IDs for the given sequences. - These are used to identify the corresponding sequence for a - *k-mer* match. - By default the IDs are counted from *0* to *m*. - ignore_masks : sized iterable object of (ndarray, dtype=bool), length=m, optional - Sequence positions to ignore. - *k-mers* that involve these sequence positions are not added - to the table. - This is used e.g. to skip repeat regions. - If provided, the list must contain one boolean mask - (or ``None``) for each sequence, and each bolean mask must - have the same length as the sequence. - By default, no sequence position is ignored. - alphabet : Alphabet, optional - The alphabet to use for this table. - It must extend the alphabets of the input `sequences`. - By default, an appropriate alphabet is inferred from the - input `sequences`. - This option is usually used for compatibility with another - sequence/table in the matching step. - spacing : None or str or list or ndarray, dtype=int, shape=(k,) - If provided, spaced *k-mers* are used instead of continuous - ones. - The value contains the *informative* positions relative to - the start of the *k-mer*, also called the *model*. - The number of *informative* positions must equal *k*. - Refer to :class:`KmerAlphabet` for more details. - - See Also - -------- - from_kmers : The same functionality based on already created *k-mers* - - Returns - ------- - table : KmerTable - The newly created table. - - Examples - -------- - - >>> sequences = [NucleotideSequence("TTATA"), NucleotideSequence("CTAG")] - >>> table = KmerTable.from_sequences( - ... 2, sequences, ref_ids=[100, 101] - ... ) - >>> print(table) - AG: (101, 2) - AT: (100, 2) - CT: (101, 0) - TA: (100, 1), (100, 3), (101, 1) - TT: (100, 0) - - Give an explicit compatible alphabet: - - >>> table = KmerTable.from_sequences( - ... 2, sequences, ref_ids=[100, 101], - ... alphabet=NucleotideSequence.ambiguous_alphabet() - ... ) - - Ignore all ``N`` in a sequence: - - >>> sequence = NucleotideSequence("ACCNTANNG") - >>> table = KmerTable.from_sequences( - ... 2, [sequence], ignore_masks=[sequence.symbols == "N"] - ... ) - >>> print(table) - AC: (0, 0) - CC: (0, 1) - TA: (0, 4) - """ - ref_ids = _compute_ref_ids(ref_ids, sequences) - ignore_masks = _compute_masks(ignore_masks, sequences) - alphabet = _compute_alphabet( - alphabet, (sequence.alphabet for sequence in sequences) - ) - - table = KmerTable(KmerAlphabet(alphabet, k, spacing)) - - # Calculate k-mers - kmers_list = [ - table._kmer_alph.create_kmers(sequence.code) - for sequence in sequences - ] - - masks = [ - _prepare_mask(table._kmer_alph, ignore_mask, len(sequence)) - for sequence, ignore_mask in zip(sequences, ignore_masks) - ] - - # Count the number of appearances of each k-mer and store the - # result in the pointer array, that is now used as count array - for kmers, mask in zip(kmers_list, masks): - table._count_masked_kmers(kmers, mask) - - # Transfrom count array into pointer array with C-array of - # appropriate size - _init_c_arrays(table._ptr_array, EntrySize.NO_BUCKETS) - - # Fill the C-arrays with the k-mer positions - for kmers, ref_id, mask in zip(kmers_list, ref_ids, masks): - table._add_kmers(kmers, ref_id, mask) - - return table - - - @staticmethod - def from_kmers(kmer_alphabet, kmers, ref_ids=None, masks=None): - """ - from_kmers(kmer_alphabet, kmers, ref_ids=None, masks=None) - - Create a :class:`KmerTable` by storing the positions of all - input *k-mers*. - - Parameters - ---------- - kmer_alphabet : KmerAlphabet - The :class:`KmerAlphabet` to use for the new table. - Should be the same alphabet that was used to calculate the - input *kmers*. - kmers : sized iterable object of (ndarray, dtype=np.int64), length=m - List where each array contains the *k-mer* codes from a - sequence. - For each array the index of the *k-mer* code in the array - is stored in the table as sequence position. - ref_ids : sized iterable object of int, length=m, optional - The reference IDs for the sequences. - These are used to identify the corresponding sequence for a - *k-mer* match. - By default the IDs are counted from *0* to *m*. - masks : sized iterable object of (ndarray, dtype=bool), length=m, optional - A *k-mer* code at a position, where the corresponding mask - is false, is not added to the table. - By default, all positions are added. - - See Also - -------- - from_sequences : The same functionality based on undecomposed sequences - - Returns - ------- - table : KmerTable - The newly created table. - - Examples - -------- - - >>> sequences = [ProteinSequence("BIQTITE"), ProteinSequence("NIQBITE")] - >>> kmer_alphabet = KmerAlphabet(ProteinSequence.alphabet, 3) - >>> kmer_codes = [kmer_alphabet.create_kmers(s.code) for s in sequences] - >>> for code in kmer_codes: - ... print(code) - [11701 4360 7879 9400 4419] - [ 6517 4364 7975 11704 4419] - >>> table = KmerTable.from_kmers( - ... kmer_alphabet, kmer_codes - ... ) - >>> print(table) - IQT: (0, 1) - IQB: (1, 1) - ITE: (0, 4), (1, 4) - NIQ: (1, 0) - QTI: (0, 2) - QBI: (1, 2) - TIT: (0, 3) - BIQ: (0, 0) - BIT: (1, 3) - """ - _check_kmer_alphabet(kmer_alphabet) - _check_multiple_kmer_bounds(kmers, kmer_alphabet) - - ref_ids = _compute_ref_ids(ref_ids, kmers) - masks = _compute_masks(masks, kmers) - - table = KmerTable(kmer_alphabet) - - masks = [ - np.ones(len(arr), dtype=np.uint8) if mask is None - # Convert boolean mask into uint8 array to be able - # to handle it as memory view - else np.frombuffer( - mask.astype(bool, copy=False), dtype=np.uint8 - ) - for mask, arr in zip(masks, kmers) - ] - - for arr, mask in zip(kmers, masks): - table._count_masked_kmers(arr, mask) - - _init_c_arrays(table._ptr_array, EntrySize.NO_BUCKETS) - - for arr, ref_id, mask in zip(kmers, ref_ids, masks): - table._add_kmers(arr, ref_id, mask) - - return table - - - @staticmethod - def from_kmer_selection(kmer_alphabet, positions, kmers, ref_ids=None): - """ - from_kmer_selection(kmer_alphabet, positions, kmers, ref_ids=None) - - Create a :class:`KmerTable` by storing the positions of a - filtered subset of input *k-mers*. - - This can be used to reduce the number of stored *k-mers* using - a *k-mer* subset selector such as :class:`MinimizerSelector`. - - Parameters - ---------- - kmer_alphabet : KmerAlphabet - The :class:`KmerAlphabet` to use for the new table. - Should be the same alphabet that was used to calculate the - input *kmers*. - positions : sized iterable object of (ndarray, shape=(n,), dtype=uint32), length=m - List where each array contains the sequence positions of - the filtered subset of *k-mers* given in `kmers`. - The list may contain multiple elements for multiple - sequences. - kmers : sized iterable object of (ndarray, shape=(n,), dtype=np.int64), length=m - List where each array contains the filtered subset of - *k-mer* codes from a sequence. - For each array the index of the *k-mer* code in the array, - is stored in the table as sequence position. - The list may contain multiple elements for multiple - sequences. - ref_ids : sized iterable object of int, length=m, optional - The reference IDs for the sequences. - These are used to identify the corresponding sequence for a - *k-mer* match. - By default the IDs are counted from *0* to *m*. - - Returns - ------- - table : KmerTable - The newly created table. - - Examples - -------- - - Reduce the size of sequence data in the table using minimizers: - - >>> sequence1 = ProteinSequence("THIS*IS*A*SEQVENCE") - >>> kmer_alph = KmerAlphabet(sequence1.alphabet, k=3) - >>> minimizer = MinimizerSelector(kmer_alph, window=4) - >>> minimizer_pos, minimizers = minimizer.select(sequence1) - >>> kmer_table = KmerTable.from_kmer_selection( - ... kmer_alph, [minimizer_pos], [minimizers] - ... ) - - Use the same :class:`MinimizerSelector` to select the minimizers - from the query sequence and match them against the table. - Although the amount of *k-mers* is reduced, matching is still - guanrateed to work, if the two sequences share identity in the - given window: - - >>> sequence2 = ProteinSequence("ANQTHER*SEQVENCE") - >>> minimizer_pos, minimizers = minimizer.select(sequence2) - >>> matches = kmer_table.match_kmer_selection(minimizer_pos, minimizers) - >>> print(matches) - [[ 9 0 11] - [12 0 14]] - >>> for query_pos, _, db_pos in matches: - ... print(sequence1) - ... print(" " * (db_pos-1) + "^" * kmer_table.k) - ... print(sequence2) - ... print(" " * (query_pos-1) + "^" * kmer_table.k) - ... print() - THIS*IS*A*SEQVENCE - ^^^ - ANQTHER*SEQVENCE - ^^^ - - THIS*IS*A*SEQVENCE - ^^^ - ANQTHER*SEQVENCE - ^^^ - - """ - _check_kmer_alphabet(kmer_alphabet) - _check_multiple_kmer_bounds(kmers, kmer_alphabet) - _check_position_shape(positions, kmers) - - ref_ids = _compute_ref_ids(ref_ids, kmers) - - table = KmerTable(kmer_alphabet) - - for arr in kmers: - table._count_kmers(arr) - - _init_c_arrays(table._ptr_array, EntrySize.NO_BUCKETS) - - for pos, arr, ref_id in zip(positions, kmers, ref_ids): - table._add_kmer_selection( - pos.astype(np.uint32, copy=False), arr, ref_id - ) - - return table - - - @staticmethod - def from_tables(tables): - """ - from_tables(tables) - - Create a :class:`KmerTable` by merging the *k-mer* positions - from existing `tables`. - - Parameters - ---------- - tables : iterable object of KmerTable - The tables to be merged. - All tables must have equal :class:`KmerAlphabet` objects, - i.e. the same *k* and equal base alphabets. - - Returns - ------- - table : KmerTable - The newly created table. - - Examples - -------- - - >>> table1 = KmerTable.from_sequences( - ... 2, [NucleotideSequence("TTATA")], ref_ids=[100] - ... ) - >>> table2 = KmerTable.from_sequences( - ... 2, [NucleotideSequence("CTAG")], ref_ids=[101] - ... ) - >>> merged_table = KmerTable.from_tables([table1, table2]) - >>> print(merged_table) - AG: (101, 2) - AT: (100, 2) - CT: (101, 0) - TA: (100, 1), (100, 3), (101, 1) - TT: (100, 0) - """ - cdef KmerTable table - - _check_same_kmer_alphabet(tables) - - merged_table = KmerTable(tables[0].kmer_alphabet) - - # Sum the number of appearances of each k-mer from the tables - for table in tables: - # 'merged_table._ptr_array' is repurposed as count array, - # This can be safely done, because in this step the pointers - # are not initialized yet. - # This may save a lot of memory because no extra array is - # required to count the number of positions for each *k-mer* - _count_table_entries( - merged_table._ptr_array, table._ptr_array, - EntrySize.NO_BUCKETS - ) - - _init_c_arrays(merged_table._ptr_array, EntrySize.NO_BUCKETS) - - for table in tables: - _append_entries(merged_table._ptr_array, table._ptr_array) - - return merged_table - - - @cython.boundscheck(False) - @cython.wraparound(False) - @staticmethod - def from_positions(kmer_alphabet, dict kmer_positions): - """ - from_positions(kmer_alphabet, kmer_positions) - - Create a :class:`KmerTable` from *k-mer* reference IDs and - positions. - This constructor is especially useful for restoring a table - from previously serialized data. - - Parameters - ---------- - kmer_alphabet : KmerAlphabet - The :class:`KmerAlphabet` to use for the new table - kmer_positions : dict of (int -> ndarray, shape=(n,2), dtype=int) - A dictionary representing the *k-mer* reference IDs and - positions to be stored in the newly created table. - It maps a *k-mer* code to a :class:`ndarray`. - To achieve a high performance the data type ``uint32`` - is preferred for the arrays. - - Returns - ------- - table : KmerTable - The newly created table. - - Examples - -------- - - >>> sequence = ProteinSequence("BIQTITE") - >>> table = KmerTable.from_sequences(3, [sequence], ref_ids=[100]) - >>> print(table) - IQT: (100, 1) - ITE: (100, 4) - QTI: (100, 2) - TIT: (100, 3) - BIQ: (100, 0) - >>> data = {kmer: table[kmer] for kmer in table} - >>> print(data) - {4360: array([[100, 1]], dtype=uint32), 4419: array([[100, 4]], dtype=uint32), 7879: array([[100, 2]], dtype=uint32), 9400: array([[100, 3]], dtype=uint32), 11701: array([[100, 0]], dtype=uint32)} - >>> restored_table = KmerTable.from_positions(table.kmer_alphabet, data) - >>> print(restored_table) - IQT: (100, 1) - ITE: (100, 4) - QTI: (100, 2) - TIT: (100, 3) - BIQ: (100, 0) - """ - cdef int64 length - cdef uint32* kmer_ptr - cdef int64 i - cdef int64 kmer - cdef uint32[:,:] positions - - table = KmerTable(kmer_alphabet) - - cdef ptr[:] ptr_array = table._ptr_array - cdef int64 alph_length = len(kmer_alphabet) - - for kmer, position_array in kmer_positions.items(): - if kmer < 0 or kmer >= alph_length: - raise AlphabetError( - f"k-mer code {kmer} does not represent a valid k-mer" - ) - positions = position_array.astype(np.uint32, copy=False) - if positions.shape[0] == 0: - # No position to add -> jump to the next k-mer - continue - if positions.shape[1] != 2: - raise IndexError( - f"Each entry in position array has {positions.shape[1]} " - f"values, but 2 were expected" - ) - - # Plus the size of array length value (int64) - length = 2 * positions.shape[0] + 2 - kmer_ptr = malloc(length * sizeof(uint32)) - if not kmer_ptr: - raise MemoryError - ptr_array[kmer] = kmer_ptr - ( kmer_ptr)[0] = length - # Jump behind the length value - kmer_ptr += 2 - - # Add entries - for i in range(positions.shape[0]): - kmer_ptr[0] = positions[i,0] - kmer_ptr += 1 - kmer_ptr[0] = positions[i,1] - kmer_ptr += 1 - - return table - - - @cython.boundscheck(False) - @cython.wraparound(False) - def match_table(self, KmerTable table, similarity_rule=None): - """ - match_table(table, similarity_rule=None) - - Find matches between the *k-mers* in this table with the - *k-mers* in another `table`. - - This means that for each *k-mer* the cartesian product between - the positions in both tables is added to the matches. - - Parameters - ---------- - table : KmerTable - The table to be matched. - Both tables must have equal :class:`KmerAlphabet` objects, - i.e. the same *k* and equal base alphabets. - similarity_rule : SimilarityRule, optional - If this parameter is given, not only exact *k-mer* matches - are considered, but also similar ones according to the given - :class:`SimilarityRule`. - - Returns - ------- - matches : ndarray, shape=(n,4), dtype=np.uint32 - The *k-mer* matches. - Each row contains one match. Each match has the following - columns: - - 0. The reference ID of the matched sequence in the other - table - 1. The sequence position of the matched sequence in the - other table - 2. The reference ID of the matched sequence in this - table - 3. The sequence position of the matched sequence in this - table - - Notes - ----- - - There is no guaranteed order of the reference IDs or - sequence positions in the returned matches. - - Examples - -------- - - >>> sequence1 = ProteinSequence("BIQTITE") - >>> table1 = KmerTable.from_sequences(3, [sequence1], ref_ids=[100]) - >>> print(table1) - IQT: (100, 1) - ITE: (100, 4) - QTI: (100, 2) - TIT: (100, 3) - BIQ: (100, 0) - >>> sequence2 = ProteinSequence("TITANITE") - >>> table2 = KmerTable.from_sequences(3, [sequence2], ref_ids=[101]) - >>> print(table2) - ANI: (101, 3) - ITA: (101, 1) - ITE: (101, 5) - NIT: (101, 4) - TAN: (101, 2) - TIT: (101, 0) - >>> print(table1.match_table(table2)) - [[101 5 100 4] - [101 0 100 3]] - """ - cdef int INIT_SIZE = 1 - - cdef int64 kmer, sim_kmer - cdef int64 match_i - cdef int64 i, j, l - cdef int64 self_length, other_length - cdef uint32* self_kmer_ptr - cdef uint32* other_kmer_ptr - - # This variable will only be used if a similarity rule exists - cdef int64[:] similar_kmers - - # Store in new variables - # to disable repetitive initialization checks - cdef ptr[:] self_ptr_array = self._ptr_array - cdef ptr[:] other_ptr_array = table._ptr_array - - _check_same_kmer_alphabet((self, table)) - - # This array will store the match positions - # As the final number of matches is unknown, a list-like - # approach is used: - # The array is initialized with a relatively small inital size - # and every time the limit would be exceeded its size is doubled - cdef int64[:,:] matches = np.empty((INIT_SIZE, 4), dtype=np.int64) - match_i = 0 - if similarity_rule is None: - for kmer in range(self_ptr_array.shape[0]): - self_kmer_ptr = self_ptr_array[kmer] - other_kmer_ptr = other_ptr_array[kmer] - # For each k-mer create the cartesian product - if self_kmer_ptr != NULL and other_kmer_ptr != NULL: - # This kmer exists for both tables - other_length = (other_kmer_ptr)[0] - self_length = (self_kmer_ptr )[0] - for i in range(2, other_length, 2): - for j in range(2, self_length, 2): - if match_i >= matches.shape[0]: - # The 'matches' array is full - # -> double its size - matches = expand(np.asarray(matches)) - matches[match_i, 0] = other_kmer_ptr[i] - matches[match_i, 1] = other_kmer_ptr[i+1] - matches[match_i, 2] = self_kmer_ptr[j] - matches[match_i, 3] = self_kmer_ptr[j+1] - match_i += 1 - - else: - for kmer in range(self_ptr_array.shape[0]): - other_kmer_ptr = other_ptr_array[kmer] - if other_kmer_ptr != NULL: - # If a similarity rule exists, iterate not only over - # the exact k-mer, but over all k-mers similar to - # the current k-mer - similar_kmers = similarity_rule.similar_kmers( - self._kmer_alph, kmer - ) - for l in range(similar_kmers.shape[0]): - sim_kmer = similar_kmers[l] - # Actual copy of the code from the other - # if-branch: - # It cannot be put properly in a cdef-function, - # as every function call would perform reference - # count changes and would decrease performance - self_kmer_ptr = self_ptr_array[sim_kmer] - if self_kmer_ptr != NULL: - other_length = (other_kmer_ptr)[0] - self_length = (self_kmer_ptr )[0] - for i in range(2, other_length, 2): - for j in range(2, self_length, 2): - if match_i >= matches.shape[0]: - matches = expand(np.asarray(matches)) - matches[match_i, 0] = other_kmer_ptr[i] - matches[match_i, 1] = other_kmer_ptr[i+1] - matches[match_i, 2] = self_kmer_ptr[j] - matches[match_i, 3] = self_kmer_ptr[j+1] - match_i += 1 - - # Trim to correct size and return - return np.asarray(matches[:match_i]) - - - @cython.boundscheck(False) - @cython.wraparound(False) - def match(self, sequence, similarity_rule=None, ignore_mask=None): - """ - match(sequence, similarity_rule=None, ignore_mask=None) - - Find matches between the *k-mers* in this table with all - overlapping *k-mers* in the given `sequence`. - *k* is determined by the table. - - Parameters - ---------- - sequence : Sequence - The sequence to be matched. - The table's base alphabet must extend the alphabet of the - sequence. - similarity_rule : SimilarityRule, optional - If this parameter is given, not only exact *k-mer* matches - are considered, but also similar ones according to the given - :class:`SimilarityRule`. - ignore_mask : ndarray, dtype=bool, optional - Boolean mask of sequence positions to ignore. - *k-mers* that involve these sequence positions are not added - to the table. - This is used e.g. to skip repeat regions. - By default, no sequence position is ignored. - - Returns - ------- - matches : ndarray, shape=(n,3), dtype=np.uint32 - The *k-mer* matches. - Each row contains one match. Each match has the following - columns: - - 0. The sequence position in the input sequence - 1. The reference ID of the matched sequence in the table - 2. The sequence position of the matched sequence in the - table - - Notes - ----- - - The matches are ordered by the first column. - - Examples - -------- - - >>> sequence1 = ProteinSequence("BIQTITE") - >>> table = KmerTable.from_sequences(3, [sequence1], ref_ids=[100]) - >>> print(table) - IQT: (100, 1) - ITE: (100, 4) - QTI: (100, 2) - TIT: (100, 3) - BIQ: (100, 0) - >>> sequence2 = ProteinSequence("TITANITE") - >>> print(table.match(sequence2)) - [[ 0 100 3] - [ 5 100 4]] - """ - cdef int INIT_SIZE = 1 - - cdef int64 kmer, sim_kmer - cdef int64 match_i - cdef int64 i, j, l - cdef int64 length - cdef uint32* kmer_ptr - - # This variable will only be used if a similarity rule exists - cdef int64[:] similar_kmers - - # Store in new variable - # to disable repetitive initialization checks - cdef ptr[:] ptr_array = self._ptr_array - - if len(sequence.code) < self._k: - raise ValueError("Sequence code is shorter than k") - if not self._kmer_alph.base_alphabet.extends(sequence.alphabet): - raise ValueError( - "The alphabet used for the k-mer index table is not equal to " - "the alphabet of the sequence" - ) - - cdef int64[:] kmers = self._kmer_alph.create_kmers(sequence.code) - cdef uint8[:] kmer_mask = _prepare_mask( - self._kmer_alph, ignore_mask, len(sequence.code) - ) - - # This array will store the match positions - # As the final number of matches is unknown, a list-like - # approach is used: - # The array is initialized with a relatively small inital size - # and every time the limit would be exceeded its size is doubled - cdef int64[:,:] matches = np.empty((INIT_SIZE, 3), dtype=np.int64) - match_i = 0 - if similarity_rule is None: - for i in range(kmers.shape[0]): - if kmer_mask[i]: - kmer = kmers[i] - kmer_ptr = ptr_array[kmer] - if kmer_ptr != NULL: - # There is at least one entry for the k-mer - length = (kmer_ptr)[0] - for j in range(2, length, 2): - if match_i >= matches.shape[0]: - # The 'matches' array is full - # -> double its size - matches = expand(np.asarray(matches)) - matches[match_i, 0] = i - matches[match_i, 1] = kmer_ptr[j] - matches[match_i, 2] = kmer_ptr[j+1] - match_i += 1 - - else: - for i in range(kmers.shape[0]): - if kmer_mask[i]: - kmer = kmers[i] - # If a similarity rule exists, iterate not only over - # the exact k-mer, but over all k-mers similar to - # the current k-mer - similar_kmers = similarity_rule.similar_kmers( - self._kmer_alph, kmer - ) - for l in range(similar_kmers.shape[0]): - sim_kmer = similar_kmers[l] - # Actual copy of the code from the other - # if-branch: - # It cannot be put properly in a cdef-function, - # as every function call would perform reference - # count changes and would decrease performance - kmer_ptr = ptr_array[sim_kmer] - if kmer_ptr != NULL: - # There is at least one entry for the k-mer - length = (kmer_ptr)[0] - for j in range(2, length, 2): - if match_i >= matches.shape[0]: - # The 'matches' array is full - # -> double its size - matches = expand(np.asarray(matches)) - matches[match_i, 0] = i - matches[match_i, 1] = kmer_ptr[j] - matches[match_i, 2] = kmer_ptr[j+1] - match_i += 1 - - # Trim to correct size and return - return np.asarray(matches[:match_i]) - - - @cython.boundscheck(False) - @cython.wraparound(False) - def match_kmer_selection(self, positions, kmers): - """ - match_kmer_selection(positions, kmers) - - Find matches between the *k-mers* in this table with the given - *k-mer* selection. - - It is intended to use this method to find matches in a table - that was created using :meth:`from_kmer_selection()`. - - Parameters - ---------- - positions : ndarray, shape=(n,), dtype=uint32 - Sequence positions of the filtered subset of *k-mers* given - in `kmers`. - kmers : ndarray, shape=(n,), dtype=np.int64 - Filtered subset of *k-mer* codes to match against. - - Returns - ------- - matches : ndarray, shape=(n,3), dtype=np.uint32 - The *k-mer* matches. - Each row contains one *k-mer* match. - Each match has the following columns: - - 0. The sequence position of the input *k-mer*, taken - from `positions` - 1. The reference ID of the matched sequence in the table - 2. The sequence position of the matched *k-mer* in the - table - - Examples - -------- - - Reduce the size of sequence data in the table using minimizers: - - >>> sequence1 = ProteinSequence("THIS*IS*A*SEQVENCE") - >>> kmer_alph = KmerAlphabet(sequence1.alphabet, k=3) - >>> minimizer = MinimizerSelector(kmer_alph, window=4) - >>> minimizer_pos, minimizers = minimizer.select(sequence1) - >>> kmer_table = KmerTable.from_kmer_selection( - ... kmer_alph, [minimizer_pos], [minimizers] - ... ) - - Use the same :class:`MinimizerSelector` to select the minimizers - from the query sequence and match them against the table. - Although the amount of *k-mers* is reduced, matching is still - guanrateed to work, if the two sequences share identity in the - given window: - - >>> sequence2 = ProteinSequence("ANQTHER*SEQVENCE") - >>> minimizer_pos, minimizers = minimizer.select(sequence2) - >>> matches = kmer_table.match_kmer_selection(minimizer_pos, minimizers) - >>> print(matches) - [[ 9 0 11] - [12 0 14]] - >>> for query_pos, _, db_pos in matches: - ... print(sequence1) - ... print(" " * (db_pos-1) + "^" * kmer_table.k) - ... print(sequence2) - ... print(" " * (query_pos-1) + "^" * kmer_table.k) - ... print() - THIS*IS*A*SEQVENCE - ^^^ - ANQTHER*SEQVENCE - ^^^ - - THIS*IS*A*SEQVENCE - ^^^ - ANQTHER*SEQVENCE - ^^^ - - """ - cdef int INIT_SIZE = 1 - - cdef int64 i, j - - cdef int64 kmer - cdef int64 match_i - cdef int64 seq_pos - cdef int64 length - cdef uint32* kmer_ptr - - # Store in new variable - # to disable repetitive initialization checks - cdef ptr[:] ptr_array = self._ptr_array - - _check_kmer_bounds(kmers, self._kmer_alph) - if positions.shape[0] != kmers.shape[0]: - raise IndexError( - f"{positions.shape[0]} positions were given " - f"for {kmers.shape[0]} k-mers" - ) - - cdef uint32[:] pos_array = positions.astype(np.uint32, copy=False) - cdef int64[:] kmer_array = kmers.astype(np.int64, copy=False) - - # This array will store the match positions - # As the final number of matches is unknown, a list-like - # approach is used: - # The array is initialized with a relatively small inital size - # and every time the limit would be exceeded its size is doubled - cdef int64[:,:] matches = np.empty((INIT_SIZE, 3), dtype=np.int64) - match_i = 0 - for i in range(kmer_array.shape[0]): - kmer = kmer_array[i] - seq_pos = pos_array[i] - kmer_ptr = ptr_array[kmer] - if kmer_ptr != NULL: - # There is at least one entry for the k-mer - length = (kmer_ptr)[0] - for j in range(2, length, 2): - if match_i >= matches.shape[0]: - # The 'matches' array is full - # -> double its size - matches = expand(np.asarray(matches)) - matches[match_i, 0] = seq_pos - matches[match_i, 1] = kmer_ptr[j] - matches[match_i, 2] = kmer_ptr[j+1] - match_i += 1 - - # Trim to correct size and return - return np.asarray(matches[:match_i]) - - - @cython.boundscheck(False) - @cython.wraparound(False) - def count(self, kmers=None): - """ - count(kmers=None) - - Count the number of occurences for each *k-mer* in the table. - - Parameters - ---------- - kmers : ndarray, dtype=np.int64, optional - The count is returned for these *k-mer* codes. - By default all *k-mers* are counted in ascending order, i.e. - ``count_for_kmer = counts[kmer]``. - - Returns - ------- - counts : ndarray, dtype=np.int64, optional - The counts for each given *k-mer*. - - Examples - -------- - >>> table = KmerTable.from_sequences( - ... k = 2, - ... sequences = [NucleotideSequence("TTATA"), NucleotideSequence("CTAG")], - ... ref_ids = [0, 1] - ... ) - >>> print(table) - AG: (1, 2) - AT: (0, 2) - CT: (1, 0) - TA: (0, 1), (0, 3), (1, 1) - TT: (0, 0) - - Count two selected *k-mers*: - - >>> print(table.count(table.kmer_alphabet.encode_multiple(["TA", "AG"]))) - [3 1] - - Count all *k-mers*: - - >>> counts = table.count() - >>> print(counts) - [0 0 1 1 0 0 0 1 0 0 0 0 3 0 0 1] - >>> for kmer, count in zip(table.kmer_alphabet.get_symbols(), counts): - ... print(kmer, count) - AA 0 - AC 0 - AG 1 - AT 1 - CA 0 - CC 0 - CG 0 - CT 1 - GA 0 - GC 0 - GG 0 - GT 0 - TA 3 - TC 0 - TG 0 - TT 1 - """ - cdef int64 i - - cdef int64 length - cdef int64 kmer - cdef int64* kmer_ptr - cdef ptr[:] ptr_array = self._ptr_array - cdef int64[:] kmer_array - cdef int64[:] counts - - if kmers is None: - counts = np.zeros(ptr_array.shape[0], dtype=np.int64) - for kmer in range(ptr_array.shape[0]): - kmer_ptr = (ptr_array[kmer]) - if kmer_ptr != NULL: - # First 64 bytes are length of C-array - length = kmer_ptr[0] - # Array length is measured in uint32 - # length = 2 * count + 2 -> rearrange formula - counts[kmer] = (length - 2) // 2 - - else: - _check_kmer_bounds(kmers, self._kmer_alph) - - kmer_array = kmers.astype(np.int64, copy=False) - counts = np.zeros(kmer_array.shape[0], dtype=np.int64) - for i in range(kmer_array.shape[0]): - kmer = kmer_array[i] - kmer_ptr = (ptr_array[kmer]) - if kmer_ptr != NULL: - length = kmer_ptr[0] - counts[i] = (length - 2) // 2 - - return np.asarray(counts) - - - @cython.boundscheck(False) - @cython.wraparound(False) - def get_kmers(self): - """ - Get the *k-mer* codes for all *k-mers* that have at least one - position in the table. - - Returns - ------- - kmers : ndarray, shape=(n,), dtype=np.int64 - The *k-mer* codes. - - Examples - -------- - - >>> sequence = ProteinSequence("BIQTITE") - >>> table = KmerTable.from_sequences(3, [sequence], ref_ids=[100]) - >>> print(table) - IQT: (100, 1) - ITE: (100, 4) - QTI: (100, 2) - TIT: (100, 3) - BIQ: (100, 0) - >>> kmer_codes = table.get_kmers() - >>> print(kmer_codes) - [ 4360 4419 7879 9400 11701] - >>> for code in kmer_codes: - ... print(table[code]) - [[100 1]] - [[100 4]] - [[100 2]] - [[100 3]] - [[100 0]] - """ - cdef int64 kmer - cdef ptr[:] ptr_array = self._ptr_array - - # Pessimistic allocation: - # The maximum number of used kmers are all possible kmers - cdef int64[:] kmers = np.zeros(ptr_array.shape[0], dtype=np.int64) - - cdef int64 i = 0 - for kmer in range(ptr_array.shape[0]): - if (ptr_array[kmer]) != NULL: - kmers[i] = kmer - i += 1 - - # Trim to correct size - return np.asarray(kmers)[:i] - - - @cython.cdivision(True) - @cython.boundscheck(False) - @cython.wraparound(False) - def __getitem__(self, int64 kmer): - cdef int64 i, j - cdef int64 length - cdef uint32* kmer_ptr - cdef uint32[:,:] positions - - if kmer >= len(self): - raise AlphabetError( - f"k-mer code {kmer} is out of bounds " - f"for the given KmerAlphabet" - ) - - kmer_ptr = self._ptr_array[kmer] - if kmer_ptr == NULL: - return np.zeros((0, 2), dtype=np.uint32) - else: - length = (kmer_ptr)[0] - positions = np.empty(((length - 2) // 2, 2), dtype=np.uint32) - i = 0 - for j in range(2, length, 2): - positions[i,0] = kmer_ptr[j] - positions[i,1] = kmer_ptr[j+1] - i += 1 - return np.asarray(positions) - - - def __len__(self): - return len(self._kmer_alph) - - - def __contains__(self, int64 kmer): - # If there is at least one entry for a k-mer, - # the pointer is not NULL - return self._ptr_array[kmer] != 0 - - - def __iter__(self): - for kmer in self.get_kmers(): - yield kmer.item() - - - def __reversed__(self): - return reversed(self.get_kmers()) - - - def __eq__(self, item): - if item is self: - return True - if type(item) != KmerTable: - return False - - # Introduce static typing to access statically typed fields - cdef KmerTable other = item - if self._kmer_alph.base_alphabet != other._kmer_alph.base_alphabet: - return False - if self._k != other._k: - return False - return _equal_c_arrays(self._ptr_array, other._ptr_array) - - - def __str__(self): - return _to_string(self) - - - def __getnewargs_ex__(self): - return (self._kmer_alph,), {} - - - def __getstate__(self): - return _pickle_c_arrays(self._ptr_array) - - - def __setstate__(self, state): - _unpickle_c_arrays(self._ptr_array, state) - - - def __dealloc__(self): - if self._is_initialized(): - _deallocate_ptrs(self._ptr_array) - - - @cython.boundscheck(False) - @cython.wraparound(False) - def _count_kmers(self, int64[:] kmers): - """ - Repurpose the pointer array as count array and add the - total number of positions for the given kmers to the values in - the count array. - - This can be safely done, because in this step the pointers are - not initialized yet. - This may save a lot of memory because no extra array is required - to count the number of positions for each *k-mer*. - """ - cdef uint32 seq_pos - cdef int64 kmer - - cdef ptr[:] count_array = self._ptr_array - - for seq_pos in range(kmers.shape[0]): - kmer = kmers[seq_pos] - count_array[kmer] += 1 - - @cython.boundscheck(False) - @cython.wraparound(False) - def _count_masked_kmers(self, int64[:] kmers, uint8[:] mask): - """ - Same as above, but with mask. - """ - cdef uint32 seq_pos - cdef int64 kmer - - cdef ptr[:] count_array = self._ptr_array - - for seq_pos in range(kmers.shape[0]): - if mask[seq_pos]: - kmer = kmers[seq_pos] - count_array[kmer] += 1 - - - @cython.boundscheck(False) - @cython.wraparound(False) - def _add_kmers(self, int64[:] kmers, uint32 ref_id, uint8[:] mask): - """ - For each *k-mer* in `kmers` add the reference ID and the - position in the array to the corresponding C-array and update - the length of the C-array. - """ - cdef uint32 seq_pos - cdef int64 current_size - cdef int64 kmer - cdef uint32* kmer_ptr - - # Store in new variable - # to disable repetitive initialization checks - cdef ptr[:] ptr_array = self._ptr_array - - if mask.shape[0] != kmers.shape[0]: - raise IndexError( - f"Mask has length {mask.shape[0]}, " - f"but there are {kmers.shape[0]} k-mers" - ) - - for seq_pos in range(kmers.shape[0]): - if mask[seq_pos]: - kmer = kmers[seq_pos] - kmer_ptr = ptr_array[kmer] - - # Append k-mer reference ID and position - current_size = ( kmer_ptr)[0] - kmer_ptr[current_size ] = ref_id - kmer_ptr[current_size + 1] = seq_pos - ( kmer_ptr)[0] = current_size + EntrySize.NO_BUCKETS - - @cython.boundscheck(False) - @cython.wraparound(False) - def _add_kmer_selection(self, uint32[:] positions, int64[:] kmers, - uint32 ref_id): - """ - For each *k-mer* in `kmers` add the reference ID and the - position from `positions` to the corresponding C-array and - update the length of the C-array. - """ - cdef uint32 i - cdef uint32 seq_pos - cdef int64 current_size - cdef int64 kmer - cdef uint32* kmer_ptr - - if positions.shape[0] != kmers.shape[0]: - raise IndexError( - f"{positions.shape[0]} positions were given " - f"for {kmers.shape[0]} k-mers" - ) - - # Store in new variable - # to disable repetitive initialization checks - cdef ptr[:] ptr_array = self._ptr_array - - for i in range(positions.shape[0]): - kmer = kmers[i] - seq_pos = positions[i] - kmer_ptr = ptr_array[kmer] - - # Append k-mer reference ID and position - current_size = ( kmer_ptr)[0] - kmer_ptr[current_size ] = ref_id - kmer_ptr[current_size + 1] = seq_pos - ( kmer_ptr)[0] = current_size + EntrySize.NO_BUCKETS - - - cdef inline bint _is_initialized(self): - # Memoryviews are not initialized on class creation - # This method checks, if the _ptr_array memoryview was - # initialized and is not None - try: - if self._ptr_array is not None: - return True - else: - return False - except AttributeError: - return False - - - - -cdef class BucketKmerTable: - """ - This class represents a *k-mer* index table. - In contrast to :class:`KmerTable` it does store each unique *k-mer* - in a separate C-array, but limits the number of C-arrays instead - to a number of buckets. - Hence, different *k-mer* may be stored in the same bucket, like in a - hash table. - This approach makes *k-mer* indices with large *k-mer* alphabets - fit into memory. - - Otherwise, the API for creating a :class:`BucketKmerTable` and - matching to it is analogous to :class:`KmerTable`. - - Attributes - ---------- - kmer_alphabet : KmerAlphabet - The internal :class:`KmerAlphabet`, that is used to - encode all overlapping *k-mers* of an input sequence. - alphabet : Alphabet - The base alphabet, from which this :class:`BucketKmerTable` was - created. - k : int - The length of the *k-mers*. - n_buckets : int - The number of buckets, the *k-mers* are divided into. - - See Also - -------- - KmerTable - - Notes - ----- - - *Memory consumption* - - For efficient mapping, a :class:`BucketKmerTable` contains a pointer - array, that contains one 64-bit pointer for each bucket. - If there is at least one position for a bucket, the corresponding - pointer points to a C-array that contains - - 1. The length of the C-array *(int64)* - 2. The *k-mers* *(int64)* - 3. The reference ID for each *k-mer* *(uint32)* - 4. The sequence position for each *k-mer* *(uint32)* - - As buckets are used, the memory requirements are limited to the number - of buckets instead of scaling with the :class:`KmerAlphabet` size. - If each bucket is used, the required memory space :math:`S` in byte - is - - .. math:: - - S = 16B + 16L - - where :math:`B` is the number of buckets and :math:`L` is the summed - length of all sequences added to the table. - - *Buckets* - - The ratio :math:`L/B` is called *load_factor*. - By default :class:`BucketKmerTable` uses a load factor of - approximately 0.8 to ensure efficient *k-mer* matching. - The number fo buckets can be adjusted by setting the - `n_buckets` parameters on :class:`BucketKmerTable` creation. - It is recommended to use :func:`bucket_number()` to compute an - appropriate number of buckets. - - *Multiprocessing* - - :class:`BucketKmerTable` objects can be used in multi-processed - setups: - Adding a large database of sequences to a table can be sped up by - splitting the database into smaller chunks and create a separate - table for each chunk in separate processes. - Eventually, the tables can be merged to one large table using - :meth:`from_tables()`. - - Since :class:`BucketKmerTable` supports the *pickle* protocol, - the matching step can also be divided into multiple processes, if - multiple sequences need to be matched. - - *Storage on hard drive* - - The most time efficient way to read/write a :class:`BucketKmerTable` - is the *pickle* format. - - *Indexing and iteration* - - Due to the higher complexity in the *k-mer* lookup compared to - :class:`KmerTable`, this class is still indexable but not iterable. - - Examples - -------- - - Create a *2-mer* index table for some nucleotide sequences: - - >>> table = BucketKmerTable.from_sequences( - ... k = 2, - ... sequences = [NucleotideSequence("TTATA"), NucleotideSequence("CTAG")], - ... ref_ids = [0, 1] - ... ) - - Display the contents of the table as - (reference ID, sequence position) tuples: - - >>> print(table) - AG: (1, 2) - AT: (0, 2) - CT: (1, 0) - TA: (0, 1), (0, 3), (1, 1) - TT: (0, 0) - - Find matches of the table with a sequence: - - >>> query = NucleotideSequence("TAG") - >>> matches = table.match(query) - >>> for query_pos, table_ref_id, table_pos in matches: - ... print("Query sequence position:", query_pos) - ... print("Table reference ID: ", table_ref_id) - ... print("Table sequence position:", table_pos) - ... print() - Query sequence position: 0 - Table reference ID: 0 - Table sequence position: 1 - - Query sequence position: 0 - Table reference ID: 0 - Table sequence position: 3 - - Query sequence position: 0 - Table reference ID: 1 - Table sequence position: 1 - - Query sequence position: 1 - Table reference ID: 1 - Table sequence position: 2 - - - Get all reference IDs and positions for a given *k-mer*: - - >>> kmer_code = table.kmer_alphabet.encode("TA") - >>> print(table[kmer_code]) - [[0 1] - [0 3] - [1 1]] - """ - - cdef object _kmer_alph - cdef int _k - cdef int64 _n_buckets - - # The pointer array is the core of the index table: - # It maps each possible k-mer bucket (represented by its code) to a - # C-array of indices. - # Each entry in a C-array contains the k-mer code, a reference ID - # and the location in that sequence where that k-mer appears - # The memory layout of each C-array is as following: - # - # (Array length) (k-mer 0) (RefID 0) (Position 0) (k-mer 1) ... - # -----int64----|--int64--|---uint32---|---uint32---|--int64-- - # - # The array length is based on 32 bit units. - # If there is no entry for a k-mer bucket, the respective pointer is - # NULL. - cdef ptr[:] _ptr_array - - - def __cinit__(self, n_buckets, kmer_alphabet): - # This check is necessary for proper memory management - # of the allocated arrays - if self._is_initialized(): - raise Exception("Duplicate call of constructor") - - self._kmer_alph = kmer_alphabet - self._k = kmer_alphabet.k - if len(self._kmer_alph) < n_buckets: - self._n_buckets = len(self._kmer_alph) - else: - self._n_buckets = n_buckets - self._ptr_array = np.zeros(self._n_buckets, dtype=np.uint64) - - - @property - def kmer_alphabet(self): - return self._kmer_alph - - @property - def alphabet(self): - return self._kmer_alph.base_alphabet - - @property - def k(self): - return self._k - - @property - def n_buckets(self): - return self._n_buckets - - @staticmethod - def from_sequences(k, sequences, ref_ids=None, ignore_masks=None, - alphabet=None, spacing=None, n_buckets=None): - """ - from_sequences(k, sequences, ref_ids=None, ignore_masks=None, - alphabet=None, spacing=None, n_buckets=None) - - Create a :class:`BucketKmerTable` by storing the positions of - all overlapping *k-mers* from the input `sequences`. - - Parameters - ---------- - k : int - The length of the *k-mers*. - sequences : sized iterable object of Sequence, length=m - The sequences to get the *k-mer* positions from. - These sequences must have equal alphabets, or one of these - sequences must have an alphabet that extends the alphabets - of all other sequences. - ref_ids : sized iterable object of int, length=m, optional - The reference IDs for the given sequences. - These are used to identify the corresponding sequence for a - *k-mer* match. - By default the IDs are counted from *0* to *m*. - ignore_masks : sized iterable object of (ndarray, dtype=bool), length=m, optional - Sequence positions to ignore. - *k-mers* that involve these sequence positions are not added - to the table. - This is used e.g. to skip repeat regions. - If provided, the list must contain one boolean mask - (or ``None``) for each sequence, and each bolean mask must - have the same length as the sequence. - By default, no sequence position is ignored. - alphabet : Alphabet, optional - The alphabet to use for this table. - It must extend the alphabets of the input `sequences`. - By default, an appropriate alphabet is inferred from the - input `sequences`. - This option is usually used for compatibility with another - sequence/table in the matching step. - spacing : None or str or list or ndarray, dtype=int, shape=(k,) - If provided, spaced *k-mers* are used instead of continuous - ones. - The value contains the *informative* positions relative to - the start of the *k-mer*, also called the *model*. - The number of *informative* positions must equal *k*. - Refer to :class:`KmerAlphabet` for more details. - n_buckets : int, optional - Set the number of buckets in the table, e.g. to use a - different load factor. - It is recommended to use :func:`bucket_number()` for this - purpose. - By default, a load factor of approximately 0.8 is used. - - See Also - -------- - from_kmers : The same functionality based on already created *k-mers* - - Returns - ------- - table : BucketKmerTable - The newly created table. - - Examples - -------- - - >>> sequences = [NucleotideSequence("TTATA"), NucleotideSequence("CTAG")] - >>> table = BucketKmerTable.from_sequences( - ... 2, sequences, ref_ids=[100, 101] - ... ) - >>> print(table) - AG: (101, 2) - AT: (100, 2) - CT: (101, 0) - TA: (100, 1), (100, 3), (101, 1) - TT: (100, 0) - - Give an explicit compatible alphabet: - - >>> table = BucketKmerTable.from_sequences( - ... 2, sequences, ref_ids=[100, 101], - ... alphabet=NucleotideSequence.ambiguous_alphabet() - ... ) - - Ignore all ``N`` in a sequence: - - >>> sequence = NucleotideSequence("ACCNTANNG") - >>> table = BucketKmerTable.from_sequences( - ... 2, [sequence], ignore_masks=[sequence.symbols == "N"] - ... ) - >>> print(table) - AC: (0, 0) - CC: (0, 1) - TA: (0, 4) - """ - ref_ids = _compute_ref_ids(ref_ids, sequences) - ignore_masks = _compute_masks(ignore_masks, sequences) - alphabet = _compute_alphabet( - alphabet, (sequence.alphabet for sequence in sequences) - ) - kmer_alphabet = KmerAlphabet(alphabet, k, spacing) - - # Calculate k-mers - kmers_list = [ - kmer_alphabet.create_kmers(sequence.code) - for sequence in sequences - ] - - if n_buckets is None: - n_kmers = np.sum([len(kmers) for kmers in kmers_list]) - n_buckets = bucket_number(n_kmers) - - table = BucketKmerTable(n_buckets, kmer_alphabet) - - masks = [ - _prepare_mask(kmer_alphabet, ignore_mask, len(sequence)) - for sequence, ignore_mask in zip(sequences, ignore_masks) - ] - - # Count the number of appearances of each k-mer and store the - # result in the pointer array, that is now used as count array - for kmers, mask in zip(kmers_list, masks): - table._count_masked_kmers(kmers, mask) - - # Transfrom count array into pointer array with C-array of - # appropriate size - _init_c_arrays(table._ptr_array, EntrySize.BUCKETS) - - # Fill the C-arrays with the k-mer positions - for kmers, ref_id, mask in zip(kmers_list, ref_ids, masks): - table._add_kmers(kmers, ref_id, mask) - - return table - - - @staticmethod - def from_kmers(kmer_alphabet, kmers, ref_ids=None, masks=None, - n_buckets=None): - """ - from_kmers(kmer_alphabet, kmers, ref_ids=None, masks=None, - n_buckets=None) - - Create a :class:`BucketKmerTable` by storing the positions of - all input *k-mers*. - - Parameters - ---------- - kmer_alphabet : KmerAlphabet - The :class:`KmerAlphabet` to use for the new table. - Should be the same alphabet that was used to calculate the - input *kmers*. - kmers : sized iterable object of (ndarray, dtype=np.int64), length=m - List where each array contains the *k-mer* codes from a - sequence. - For each array the index of the *k-mer* code in the array - is stored in the table as sequence position. - ref_ids : sized iterable object of int, length=m, optional - The reference IDs for the sequences. - These are used to identify the corresponding sequence for a - *k-mer* match. - By default the IDs are counted from *0* to *m*. - masks : sized iterable object of (ndarray, dtype=bool), length=m, optional - A *k-mer* code at a position, where the corresponding mask - is false, is not added to the table. - By default, all positions are added. - n_buckets : int, optional - Set the number of buckets in the table, e.g. to use a - different load factor. - It is recommended to use :func:`bucket_number()` for this - purpose. - By default, a load factor of approximately 0.8 is used. - - See Also - -------- - from_sequences : The same functionality based on undecomposed sequences - - Returns - ------- - table : BucketKmerTable - The newly created table. - - Examples - -------- - - >>> sequences = [ProteinSequence("BIQTITE"), ProteinSequence("NIQBITE")] - >>> kmer_alphabet = KmerAlphabet(ProteinSequence.alphabet, 3) - >>> kmer_codes = [kmer_alphabet.create_kmers(s.code) for s in sequences] - >>> for code in kmer_codes: - ... print(code) - [11701 4360 7879 9400 4419] - [ 6517 4364 7975 11704 4419] - >>> table = BucketKmerTable.from_kmers(kmer_alphabet, kmer_codes) - >>> print(table) - IQT: (0, 1) - IQB: (1, 1) - ITE: (0, 4), (1, 4) - NIQ: (1, 0) - QTI: (0, 2) - QBI: (1, 2) - TIT: (0, 3) - BIQ: (0, 0) - BIT: (1, 3) - """ - _check_kmer_alphabet(kmer_alphabet) - _check_multiple_kmer_bounds(kmers, kmer_alphabet) - - ref_ids = _compute_ref_ids(ref_ids, kmers) - masks = _compute_masks(masks, kmers) - - if n_buckets is None: - n_kmers = np.sum([len(e) for e in kmers]) - n_buckets = bucket_number(n_kmers) - - table = BucketKmerTable(n_buckets, kmer_alphabet) - - masks = [ - np.ones(len(arr), dtype=np.uint8) if mask is None - # Convert boolean mask into uint8 array to be able - # to handle it as memory view - else np.frombuffer( - mask.astype(bool, copy=False), dtype=np.uint8 - ) - for mask, arr in zip(masks, kmers) - ] - - for arr, mask in zip(kmers, masks): - table._count_masked_kmers(arr, mask) - - _init_c_arrays(table._ptr_array, EntrySize.BUCKETS) - - for arr, ref_id, mask in zip(kmers, ref_ids, masks): - table._add_kmers(arr, ref_id, mask) - - return table - - - @staticmethod - def from_kmer_selection(kmer_alphabet, positions, kmers, ref_ids=None, - n_buckets=None): - """ - from_kmer_selection(kmer_alphabet, positions, kmers, ref_ids=None, - n_buckets=None) - - Create a :class:`BucketKmerTable` by storing the positions of a - filtered subset of input *k-mers*. - - This can be used to reduce the number of stored *k-mers* using - a *k-mer* subset selector such as :class:`MinimizerSelector`. - - Parameters - ---------- - kmer_alphabet : KmerAlphabet - The :class:`KmerAlphabet` to use for the new table. - Should be the same alphabet that was used to calculate the - input *kmers*. - positions : sized iterable object of (ndarray, shape=(n,), dtype=uint32), length=m - List where each array contains the sequence positions of - the filtered subset of *k-mers* given in `kmers`. - The list may contain multiple elements for multiple - sequences. - kmers : sized iterable object of (ndarray, shape=(n,), dtype=np.int64), length=m - List where each array contains the filtered subset of - *k-mer* codes from a sequence. - For each array the index of the *k-mer* code in the array, - is stored in the table as sequence position. - The list may contain multiple elements for multiple - sequences. - ref_ids : sized iterable object of int, length=m, optional - The reference IDs for the sequences. - These are used to identify the corresponding sequence for a - *k-mer* match. - By default the IDs are counted from *0* to *m*. - n_buckets : int, optional - Set the number of buckets in the table, e.g. to use a - different load factor. - It is recommended to use :func:`bucket_number()` for this - purpose. - By default, a load factor of approximately 0.8 is used. - - Returns - ------- - table : BucketKmerTable - The newly created table. - - Examples - -------- - - Reduce the size of sequence data in the table using minimizers: - - >>> sequence1 = ProteinSequence("THIS*IS*A*SEQVENCE") - >>> kmer_alph = KmerAlphabet(sequence1.alphabet, k=3) - >>> minimizer = MinimizerSelector(kmer_alph, window=4) - >>> minimizer_pos, minimizers = minimizer.select(sequence1) - >>> kmer_table = BucketKmerTable.from_kmer_selection( - ... kmer_alph, [minimizer_pos], [minimizers] - ... ) - - Use the same :class:`MinimizerSelector` to select the minimizers - from the query sequence and match them against the table. - Although the amount of *k-mers* is reduced, matching is still - guanrateed to work, if the two sequences share identity in the - given window: - - >>> sequence2 = ProteinSequence("ANQTHER*SEQVENCE") - >>> minimizer_pos, minimizers = minimizer.select(sequence2) - >>> matches = kmer_table.match_kmer_selection(minimizer_pos, minimizers) - >>> print(matches) - [[ 9 0 11] - [12 0 14]] - >>> for query_pos, _, db_pos in matches: - ... print(sequence1) - ... print(" " * (db_pos-1) + "^" * kmer_table.k) - ... print(sequence2) - ... print(" " * (query_pos-1) + "^" * kmer_table.k) - ... print() - THIS*IS*A*SEQVENCE - ^^^ - ANQTHER*SEQVENCE - ^^^ - - THIS*IS*A*SEQVENCE - ^^^ - ANQTHER*SEQVENCE - ^^^ - - """ - _check_kmer_alphabet(kmer_alphabet) - _check_multiple_kmer_bounds(kmers, kmer_alphabet) - _check_position_shape(positions, kmers) - - ref_ids = _compute_ref_ids(ref_ids, kmers) - - if n_buckets is None: - n_kmers = np.sum([len(e) for e in kmers]) - n_buckets = bucket_number(n_kmers) - - table = BucketKmerTable(n_buckets, kmer_alphabet) - - for arr in kmers: - table._count_kmers(arr) - - _init_c_arrays(table._ptr_array, EntrySize.BUCKETS) - - for pos, arr, ref_id in zip(positions, kmers, ref_ids): - table._add_kmer_selection( - pos.astype(np.uint32, copy=False), arr, ref_id - ) - - return table - - - @staticmethod - def from_tables(tables): - """ - from_tables(tables) - - Create a :class:`BucketKmerTable` by merging the *k-mer* - positions from existing `tables`. - - Parameters - ---------- - tables : iterable object of BucketKmerTable - The tables to be merged. - All tables must have equal number of buckets and equal - :class:`KmerAlphabet` objects, i.e. the same *k* and equal - base alphabets. - - Returns - ------- - table : BucketKmerTable - The newly created table. - - Examples - -------- - To ensure that all tables have the same number of buckets, - `n_buckets` need to be set on table creation. - - >>> # The sequence length is not exactly the length of resulting k-mers, - >>> # but it is close enough for bucket computation - >>> n_buckets = bucket_number(len("TTATA") + len("CTAG")) - >>> table1 = BucketKmerTable.from_sequences( - ... 2, [NucleotideSequence("TTATA")], ref_ids=[100], - ... n_buckets=n_buckets - ... ) - >>> table2 = BucketKmerTable.from_sequences( - ... 2, [NucleotideSequence("CTAG")], ref_ids=[101], - ... n_buckets=n_buckets - ... ) - >>> merged_table = BucketKmerTable.from_tables([table1, table2]) - >>> print(merged_table) - AG: (101, 2) - AT: (100, 2) - CT: (101, 0) - TA: (100, 1), (100, 3), (101, 1) - TT: (100, 0) - """ - cdef BucketKmerTable table - - _check_same_kmer_alphabet(tables) - _check_same_buckets(tables) - - merged_table = BucketKmerTable( - tables[0].n_buckets, - tables[0].kmer_alphabet - ) - - # Sum the number of appearances of each k-mer from the tables - for table in tables: - _count_table_entries( - merged_table._ptr_array, table._ptr_array, - EntrySize.BUCKETS - ) - - _init_c_arrays(merged_table._ptr_array, EntrySize.BUCKETS) - - for table in tables: - _append_entries(merged_table._ptr_array, table._ptr_array) - - return merged_table - - - @cython.cdivision(True) - @cython.boundscheck(False) - @cython.wraparound(False) - def match_table(self, BucketKmerTable table, similarity_rule=None): - """ - match_table(table, similarity_rule=None) - - Find matches between the *k-mers* in this table with the - *k-mers* in another `table`. - - This means that for each *k-mer* the cartesian product between - the positions in both tables is added to the matches. - - Parameters - ---------- - table : BucketKmerTable - The table to be matched. - Both tables must have equal number of buckets and equal - :class:`KmerAlphabet` objects, i.e. the same *k* and equal - base alphabets. - similarity_rule : SimilarityRule, optional - If this parameter is given, not only exact *k-mer* matches - are considered, but also similar ones according to the given - :class:`SimilarityRule`. - - Returns - ------- - matches : ndarray, shape=(n,4), dtype=np.uint32 - The *k-mer* matches. - Each row contains one match. Each match has the following - columns: - - 0. The reference ID of the matched sequence in the other - table - 1. The sequence position of the matched sequence in the - other table - 2. The reference ID of the matched sequence in this - table - 3. The sequence position of the matched sequence in this - table - - Notes - ----- - - - There is no guaranteed order of the reference IDs or - sequence positions in the returned matches. - - Examples - -------- - To ensure that both tables have the same number of buckets, - `n_buckets` need to be set on table creation. - - >>> # The sequence length is not exactly the length of resulting k-mers, - >>> # but it is close enouggh for bucket computation - >>> n_buckets = bucket_number(max(len("BIQTITE"), len("TITANITE"))) - >>> sequence1 = ProteinSequence("BIQTITE") - >>> table1 = BucketKmerTable.from_sequences(3, [sequence1], ref_ids=[100]) - >>> print(table1) - IQT: (100, 1) - ITE: (100, 4) - QTI: (100, 2) - TIT: (100, 3) - BIQ: (100, 0) - >>> sequence2 = ProteinSequence("TITANITE") - >>> table2 = BucketKmerTable.from_sequences(3, [sequence2], ref_ids=[101]) - >>> print(table2) - ANI: (101, 3) - ITA: (101, 1) - ITE: (101, 5) - NIT: (101, 4) - TAN: (101, 2) - TIT: (101, 0) - >>> print(table1.match_table(table2)) - [[101 0 100 3] - [101 5 100 4]] - """ - cdef int INIT_SIZE = 1 - - cdef int64 bucket, sim_bucket - cdef int64 self_kmer, other_kmer, sim_kmer - cdef int64 match_i - cdef int64 i, j, l - cdef int64 self_length, other_length - cdef uint32* self_bucket_ptr - cdef uint32* other_bucket_ptr - - # This variable will only be used if a similarity rule exists - cdef int64[:] similar_kmers - - # Store in new variables - # to disable repetitive initialization checks - cdef ptr[:] self_ptr_array = self._ptr_array - cdef ptr[:] other_ptr_array = table._ptr_array - - _check_same_kmer_alphabet((self, table)) - _check_same_buckets((self, table)) - - # This array will store the match positions - # As the final number of matches is unknown, a list-like - # approach is used: - # The array is initialized with a relatively small inital size - # and every time the limit would be exceeded its size is doubled - cdef int64[:,:] matches = np.empty((INIT_SIZE, 4), dtype=np.int64) - match_i = 0 - if similarity_rule is None: - for bucket in range(self_ptr_array.shape[0]): - self_bucket_ptr = self_ptr_array[bucket] - other_bucket_ptr = other_ptr_array[bucket] - if self_bucket_ptr != NULL and other_bucket_ptr != NULL: - # This bucket exists for both tables - other_length = (other_bucket_ptr)[0] - self_length = (self_bucket_ptr )[0] - for i in range(2, other_length, 4): - # Hacky syntax to achieve casting to int64* - # after offset is applied - other_kmer = ((other_bucket_ptr + i))[0] - for j in range(2, self_length, 4): - self_kmer = ((self_bucket_ptr + j))[0] - if self_kmer == other_kmer: - # The k-mers are not only in the same - # bucket, but they are actually equal - if match_i >= matches.shape[0]: - # The 'matches' array is full - # -> double its size - matches = expand(np.asarray(matches)) - matches[match_i, 0] = other_bucket_ptr[i+2] - matches[match_i, 1] = other_bucket_ptr[i+3] - matches[match_i, 2] = self_bucket_ptr[j+2] - matches[match_i, 3] = self_bucket_ptr[j+3] - match_i += 1 - - else: - for bucket in range(self_ptr_array.shape[0]): - other_bucket_ptr = other_ptr_array[bucket] - if other_bucket_ptr != NULL: - other_length = (other_bucket_ptr)[0] - for i in range(2, other_length, 4): - other_kmer = ((other_bucket_ptr + i))[0] - # If a similarity rule exists, iterate not only over - # the exact k-mer, but over all k-mers similar to - # the current k-mer - similar_kmers = similarity_rule.similar_kmers( - self._kmer_alph, other_kmer - ) - for l in range(similar_kmers.shape[0]): - sim_kmer = similar_kmers[l] - sim_bucket = sim_kmer % self._n_buckets - self_bucket_ptr = self_ptr_array[sim_bucket] - if self_bucket_ptr != NULL: - self_length = (self_bucket_ptr)[0] - for j in range(2, self_length, 4): - self_kmer = ((self_bucket_ptr + j))[0] - if self_kmer == sim_kmer: - if match_i >= matches.shape[0]: - # The 'matches' array is full - # -> double its size - matches = expand(np.asarray(matches)) - matches[match_i, 0] = other_bucket_ptr[i+2] - matches[match_i, 1] = other_bucket_ptr[i+3] - matches[match_i, 2] = self_bucket_ptr[j+2] - matches[match_i, 3] = self_bucket_ptr[j+3] - match_i += 1 - - # Trim to correct size and return - return np.asarray(matches[:match_i]) - - - @cython.cdivision(True) - @cython.boundscheck(False) - @cython.wraparound(False) - def match(self, sequence, similarity_rule=None, ignore_mask=None): - """ - match(sequence, similarity_rule=None, ignore_mask=None) - - Find matches between the *k-mers* in this table with all - overlapping *k-mers* in the given `sequence`. - *k* is determined by the table. - - Parameters - ---------- - sequence : Sequence - The sequence to be matched. - The table's base alphabet must extend the alphabet of the - sequence. - similarity_rule : SimilarityRule, optional - If this parameter is given, not only exact *k-mer* matches - are considered, but also similar ones according to the given - :class:`SimilarityRule`. - ignore_mask : ndarray, dtype=bool, optional - Boolean mask of sequence positions to ignore. - *k-mers* that involve these sequence positions are not added - to the table. - This is used e.g. to skip repeat regions. - By default, no sequence position is ignored. - - Returns - ------- - matches : ndarray, shape=(n,3), dtype=np.uint32 - The *k-mer* matches. - Each row contains one match. Each match has the following - columns: - - 0. The sequence position in the input sequence - 1. The reference ID of the matched sequence in the table - 2. The sequence position of the matched sequence in the - table - - Notes - ----- - - The matches are ordered by the first column. - - Examples - -------- - - >>> sequence1 = ProteinSequence("BIQTITE") - >>> table = BucketKmerTable.from_sequences(3, [sequence1], ref_ids=[100]) - >>> print(table) - IQT: (100, 1) - ITE: (100, 4) - QTI: (100, 2) - TIT: (100, 3) - BIQ: (100, 0) - >>> sequence2 = ProteinSequence("TITANITE") - >>> print(table.match(sequence2)) - [[ 0 100 3] - [ 5 100 4]] - """ - cdef int INIT_SIZE = 1 - - cdef int64 bucket - cdef int64 self_kmer, other_kmer, sim_kmer - cdef int64 match_i - cdef int64 i, l - cdef int64 length - cdef uint32* bucket_ptr - cdef uint32* array_stop - - # This variable will only be used if a similarity rule exists - cdef int64[:] similar_kmers - - # Store in new variable - # to disable repetitive initialization checks - cdef ptr[:] ptr_array = self._ptr_array - - if len(sequence.code) < self._k: - raise ValueError("Sequence code is shorter than k") - if not self._kmer_alph.base_alphabet.extends(sequence.alphabet): - raise ValueError( - "The alphabet used for the k-mer index table is not equal to " - "the alphabet of the sequence" - ) - - cdef int64[:] kmers = self._kmer_alph.create_kmers(sequence.code) - cdef uint8[:] kmer_mask = _prepare_mask( - self._kmer_alph, ignore_mask, len(sequence.code) - ) - - # This array will store the match positions - # As the final number of matches is unknown, a list-like - # approach is used: - # The array is initialized with a relatively small inital size - # and every time the limit would be exceeded its size is doubled - cdef int64[:,:] matches = np.empty((INIT_SIZE, 3), dtype=np.int64) - match_i = 0 - if similarity_rule is None: - for i in range(kmers.shape[0]): - if kmer_mask[i]: - other_kmer = kmers[i] - bucket = other_kmer % self._n_buckets - bucket_ptr = ptr_array[bucket] - if bucket_ptr != NULL: - # There is at least one entry in this bucket - length = (bucket_ptr)[0] - array_stop = bucket_ptr + length - bucket_ptr += 2 - while bucket_ptr < array_stop: - self_kmer = (bucket_ptr)[0] - if self_kmer == other_kmer: - # The k-mers are not only in the same - # bucket, but they are actually equal - if match_i >= matches.shape[0]: - # The 'matches' array is full - # -> double its size - matches = expand(np.asarray(matches)) - matches[match_i, 0] = i - bucket_ptr += 2 - matches[match_i, 1] = bucket_ptr[0] - bucket_ptr += 1 - matches[match_i, 2] = bucket_ptr[0] - bucket_ptr += 1 - match_i += 1 - else: - bucket_ptr += EntrySize.BUCKETS - - else: - for i in range(kmers.shape[0]): - if kmer_mask[i]: - other_kmer = kmers[i] - # If a similarity rule exists, iterate not only over - # the exact k-mer, but over all k-mers similar to - # the current k-mer - similar_kmers = similarity_rule.similar_kmers( - self._kmer_alph, other_kmer - ) - for l in range(similar_kmers.shape[0]): - sim_kmer = similar_kmers[l] - bucket = sim_kmer % self._n_buckets - # Actual copy of the code from the other - # if-branch: - # It cannot be put properly in a cdef-function, - # as every function call would perform reference - # count changes and would decrease performance - bucket_ptr = ptr_array[bucket] - if bucket_ptr != NULL: - # There is at least one entry in this bucket - length = (bucket_ptr)[0] - array_stop = bucket_ptr + length - bucket_ptr += 2 - while bucket_ptr < array_stop: - self_kmer = (bucket_ptr)[0] - if self_kmer == sim_kmer: - # The k-mers are not only in the same - # bucket, but they are actually equal - if match_i >= matches.shape[0]: - # The 'matches' array is full - # -> double its size - matches = expand(np.asarray(matches)) - matches[match_i, 0] = i - bucket_ptr += 2 - matches[match_i, 1] = bucket_ptr[0] - bucket_ptr += 1 - matches[match_i, 2] = bucket_ptr[0] - bucket_ptr += 1 - match_i += 1 - else: - bucket_ptr += EntrySize.BUCKETS - - # Trim to correct size and return - return np.asarray(matches[:match_i]) - - - @cython.cdivision(True) - @cython.boundscheck(False) - @cython.wraparound(False) - def match_kmer_selection(self, positions, kmers): - """ - match_kmer_selection(positions, kmers) - - Find matches between the *k-mers* in this table with the given - *k-mer* selection. - - It is intended to use this method to find matches in a table - that was created using :meth:`from_kmer_selection()`. - - Parameters - ---------- - positions : ndarray, shape=(n,), dtype=uint32 - Sequence positions of the filtered subset of *k-mers* given - in `kmers`. - kmers : ndarray, shape=(n,), dtype=np.int64 - Filtered subset of *k-mer* codes to match against. - - Returns - ------- - matches : ndarray, shape=(n,3), dtype=np.uint32 - The *k-mer* matches. - Each row contains one *k-mer* match. - Each match has the following columns: - - 0. The sequence position of the input *k-mer*, taken - from `positions` - 1. The reference ID of the matched sequence in the table - 2. The sequence position of the matched *k-mer* in the - table - - Examples - -------- - - Reduce the size of sequence data in the table using minimizers: - - >>> sequence1 = ProteinSequence("THIS*IS*A*SEQVENCE") - >>> kmer_alph = KmerAlphabet(sequence1.alphabet, k=3) - >>> minimizer = MinimizerSelector(kmer_alph, window=4) - >>> minimizer_pos, minimizers = minimizer.select(sequence1) - >>> kmer_table = BucketKmerTable.from_kmer_selection( - ... kmer_alph, [minimizer_pos], [minimizers] - ... ) - - Use the same :class:`MinimizerSelector` to select the minimizers - from the query sequence and match them against the table. - Although the amount of *k-mers* is reduced, matching is still - guanrateed to work, if the two sequences share identity in the - given window: - - >>> sequence2 = ProteinSequence("ANQTHER*SEQVENCE") - >>> minimizer_pos, minimizers = minimizer.select(sequence2) - >>> matches = kmer_table.match_kmer_selection(minimizer_pos, minimizers) - >>> print(matches) - [[ 9 0 11] - [12 0 14]] - >>> for query_pos, _, db_pos in matches: - ... print(sequence1) - ... print(" " * (db_pos-1) + "^" * kmer_table.k) - ... print(sequence2) - ... print(" " * (query_pos-1) + "^" * kmer_table.k) - ... print() - THIS*IS*A*SEQVENCE - ^^^ - ANQTHER*SEQVENCE - ^^^ - - THIS*IS*A*SEQVENCE - ^^^ - ANQTHER*SEQVENCE - ^^^ - - """ - cdef int INIT_SIZE = 1 - - cdef int64 i - - cdef int64 bucket - cdef int64 self_kmer, other_kmer - cdef int64 match_i - cdef int64 seq_pos - cdef int64 length - cdef uint32* bucket_ptr - cdef uint32* array_stop - - # Store in new variable - # to disable repetitive initialization checks - cdef ptr[:] ptr_array = self._ptr_array - - _check_kmer_bounds(kmers, self._kmer_alph) - if positions.shape[0] != kmers.shape[0]: - raise IndexError( - f"{positions.shape[0]} positions were given " - f"for {kmers.shape[0]} k-mers" - ) - - cdef uint32[:] pos_array = positions.astype(np.uint32, copy=False) - cdef int64[:] kmer_array = kmers.astype(np.int64, copy=False) - - # This array will store the match positions - # As the final number of matches is unknown, a list-like - # approach is used: - # The array is initialized with a relatively small inital size - # and every time the limit would be exceeded its size is doubled - cdef int64[:,:] matches = np.empty((INIT_SIZE, 3), dtype=np.int64) - match_i = 0 - for i in range(kmer_array.shape[0]): - other_kmer = kmer_array[i] - seq_pos = pos_array[i] - bucket = other_kmer % self._n_buckets - bucket_ptr = ptr_array[bucket] - if bucket_ptr != NULL: - # There is at least one entry in this bucket - length = (bucket_ptr)[0] - array_stop = bucket_ptr + length - bucket_ptr += 2 - while bucket_ptr < array_stop: - self_kmer = (bucket_ptr)[0] - if self_kmer == other_kmer: - # The k-mers are not only in the same - # bucket, but they are actually equal - if match_i >= matches.shape[0]: - # The 'matches' array is full - # -> double its size - matches = expand(np.asarray(matches)) - matches[match_i, 0] = seq_pos - bucket_ptr += 2 - matches[match_i, 1] = bucket_ptr[0] - bucket_ptr += 1 - matches[match_i, 2] = bucket_ptr[0] - bucket_ptr += 1 - match_i += 1 - else: - bucket_ptr += EntrySize.BUCKETS - - # Trim to correct size and return - return np.asarray(matches[:match_i]) - - - @cython.cdivision(True) - @cython.boundscheck(False) - @cython.wraparound(False) - def count(self, kmers): - """ - count(kmers=None) - - Count the number of occurences for each *k-mer* in the table. - - Parameters - ---------- - kmers : ndarray, dtype=np.int64, optional - The count is returned for these *k-mer* codes. - By default all *k-mers* are counted in ascending order, i.e. - ``count_for_kmer = counts[kmer]``. - - Returns - ------- - counts : ndarray, dtype=np.int64, optional - The counts for each given *k-mer*. - - Notes - ----- - As each bucket need to be inspected for the actual *k-mer* - entries, this method requires far more computation time than its - :class:`KmerTable` equivalent. - - Examples - -------- - >>> table = BucketKmerTable.from_sequences( - ... k = 2, - ... sequences = [NucleotideSequence("TTATA"), NucleotideSequence("CTAG")], - ... ref_ids = [0, 1] - ... ) - >>> print(table) - AG: (1, 2) - AT: (0, 2) - CT: (1, 0) - TA: (0, 1), (0, 3), (1, 1) - TT: (0, 0) - - Count two selected *k-mers*: - - >>> print(table.count(table.kmer_alphabet.encode_multiple(["TA", "AG"]))) - [3 1] - """ - cdef int64 i - - cdef int64 bucket - cdef int64 kmer, self_kmer - cdef int64 length - cdef uint32* bucket_ptr - cdef uint32* array_stop - cdef ptr[:] ptr_array = self._ptr_array - - _check_kmer_bounds(kmers, self._kmer_alph) - cdef int64[:] kmer_array = kmers.astype(np.int64, copy=False) - cdef int64[:] counts = np.zeros(kmer_array.shape[0], dtype=np.int64) - - for i in range(kmer_array.shape[0]): - kmer = kmer_array[i] - bucket = kmer % self._n_buckets - bucket_ptr = (ptr_array[bucket]) - if bucket_ptr != NULL: - length = (bucket_ptr)[0] - array_stop = bucket_ptr + length - bucket_ptr += 2 - while bucket_ptr < array_stop: - self_kmer = (bucket_ptr)[0] - if self_kmer == kmer: - counts[i] += 1 - bucket_ptr += EntrySize.BUCKETS - - return np.asarray(counts) - - - @cython.boundscheck(False) - @cython.wraparound(False) - def get_kmers(self): - """ - Get the *k-mer* codes for all *k-mers* that have at least one - position in the table. - - Returns - ------- - kmers : ndarray, shape=(n,), dtype=np.int64 - The *k-mer* codes. - - Notes - ----- - As each bucket need to be inspected for the actual *k-mer* - entries, this method requires far more computation time than its - :class:`KmerTable` equivalent. - - Examples - -------- - - >>> sequence = ProteinSequence("BIQTITE") - >>> table = BucketKmerTable.from_sequences(3, [sequence], ref_ids=[100]) - >>> print(table) - IQT: (100, 1) - ITE: (100, 4) - QTI: (100, 2) - TIT: (100, 3) - BIQ: (100, 0) - >>> kmer_codes = table.get_kmers() - >>> print(kmer_codes) - [ 4360 4419 7879 9400 11701] - >>> for code in kmer_codes: - ... print(table[code]) - [[100 1]] - [[100 4]] - [[100 2]] - [[100 3]] - [[100 0]] - """ - cdef int64 bucket - cdef int64 kmer - cdef int64 length - cdef uint32* bucket_ptr - cdef uint32* array_stop - cdef ptr[:] ptr_array = self._ptr_array - - cdef cpp_set[int64] kmer_set - - for bucket in range(ptr_array.shape[0]): - bucket_ptr = (ptr_array[bucket]) - if bucket_ptr != NULL: - length = (bucket_ptr)[0] - array_stop = bucket_ptr + length - bucket_ptr += 2 - while bucket_ptr < array_stop: - kmer = (bucket_ptr)[0] - kmer_set.insert(kmer) - bucket_ptr += EntrySize.BUCKETS - - cdef int64[:] kmers = np.zeros(kmer_set.size(), dtype=np.int64) - cdef int64 i = 0 - for kmer in kmer_set: - kmers[i] = kmer - i += 1 - return np.sort(np.asarray(kmers)) - - - @cython.cdivision(True) - @cython.boundscheck(False) - @cython.wraparound(False) - def __getitem__(self, int64 kmer): - cdef int64 i, j - cdef int64 self_kmer - cdef int64 length - cdef uint32* bucket_ptr - cdef uint32[:,:] positions - - if kmer >= len(self): - raise AlphabetError( - f"k-mer code {kmer} is out of bounds " - f"for the given KmerAlphabet" - ) - - bucket_ptr = self._ptr_array[kmer % self._n_buckets] - if bucket_ptr == NULL: - return np.zeros((0, 2), dtype=np.uint32) - else: - length = (bucket_ptr)[0] - # Pessimistic array allocation: - # All k-mer positions in bucket belong to the requested k-mer - positions = np.empty(((length - 2) // 4, 2), dtype=np.uint32) - i = 0 - for j in range(2, length, 4): - self_kmer = bucket_ptr[j] - if self_kmer == kmer: - positions[i,0] = bucket_ptr[j+2] - positions[i,1] = bucket_ptr[j+3] - i += 1 - # Trim to correct size - return np.asarray(positions)[:i] - - - def __len__(self): - return len(self._kmer_alph) - - - def __eq__(self, item): - if item is self: - return True - if type(item) != BucketKmerTable: - return False - - # Introduce static typing to access statically typed fields - cdef BucketKmerTable other = item - if self._kmer_alph.base_alphabet != other._kmer_alph.base_alphabet: - return False - if self._k != other._k: - return False - if self._n_buckets != other._n_buckets: - return False - return _equal_c_arrays(self._ptr_array, other._ptr_array) - - - def __str__(self): - return _to_string(self) - - - def __getnewargs_ex__(self): - return (self._n_buckets, self._kmer_alph), {} - - - def __getstate__(self): - return _pickle_c_arrays(self._ptr_array) - - def __setstate__(self, state): - _unpickle_c_arrays(self._ptr_array, state) - - - def __dealloc__(self): - if self._is_initialized(): - _deallocate_ptrs(self._ptr_array) - - - ## These private methods work analogous to KmerTable - - @cython.cdivision(True) - @cython.boundscheck(False) - @cython.wraparound(False) - def _count_kmers(self, int64[:] kmers): - cdef uint32 seq_pos - cdef int64 kmer - - cdef ptr[:] count_array = self._ptr_array - - for seq_pos in range(kmers.shape[0]): - kmer = kmers[seq_pos] - # Pool all k-mers that should go into the same bucket - count_array[kmer % self._n_buckets] += 1 - - @cython.cdivision(True) - @cython.boundscheck(False) - @cython.wraparound(False) - def _count_masked_kmers(self, int64[:] kmers, uint8[:] mask): - cdef uint32 seq_pos - cdef int64 kmer - - cdef ptr[:] count_array = self._ptr_array - - for seq_pos in range(kmers.shape[0]): - if mask[seq_pos]: - kmer = kmers[seq_pos] - # Pool all k-mers that should go into the same bucket - count_array[kmer % self._n_buckets] += 1 - - - @cython.cdivision(True) - @cython.boundscheck(False) - @cython.wraparound(False) - def _add_kmers(self, int64[:] kmers, uint32 ref_id, uint8[:] mask): - cdef uint32 seq_pos - cdef int64 current_size - cdef int64 kmer - cdef uint32* bucket_ptr - cdef uint32* kmer_val_ptr - - # Store in new variable - # to disable repetitive initialization checks - cdef ptr[:] ptr_array = self._ptr_array - - if mask.shape[0] != kmers.shape[0]: - raise IndexError( - f"Mask has length {mask.shape[0]}, " - f"but there are {kmers.shape[0]} k-mers" - ) - - for seq_pos in range(kmers.shape[0]): - if mask[seq_pos]: - kmer = kmers[seq_pos] - bucket_ptr = ptr_array[kmer % self._n_buckets] - - # Append k-mer, reference ID and position - current_size = ( bucket_ptr)[0] - kmer_val_ptr = &bucket_ptr[current_size] - ( kmer_val_ptr)[0] = kmer - bucket_ptr[current_size + 2] = ref_id - bucket_ptr[current_size + 3] = seq_pos - ( bucket_ptr)[0] = current_size + EntrySize.BUCKETS - - - @cython.cdivision(True) - @cython.boundscheck(False) - @cython.wraparound(False) - def _add_kmer_selection(self, uint32[:] positions, int64[:] kmers, - uint32 ref_id): - cdef uint32 i - cdef uint32 seq_pos - cdef int64 current_size - cdef int64 kmer - cdef uint32* bucket_ptr - cdef uint32* kmer_val_ptr - - if positions.shape[0] != kmers.shape[0]: - raise IndexError( - f"{positions.shape[0]} positions were given " - f"for {kmers.shape[0]} k-mers" - ) - - # Store in new variable - # to disable repetitive initialization checks - cdef ptr[:] ptr_array = self._ptr_array - - for i in range(positions.shape[0]): - kmer = kmers[i] - seq_pos = positions[i] - bucket_ptr = ptr_array[kmer % self._n_buckets] - - # Append k-mer reference ID and position - current_size = ( bucket_ptr)[0] - kmer_val_ptr = &bucket_ptr[current_size] - ( kmer_val_ptr)[0] = kmer - bucket_ptr[current_size + 2] = ref_id - bucket_ptr[current_size + 3] = seq_pos - ( bucket_ptr)[0] = current_size + EntrySize.BUCKETS - - - cdef inline bint _is_initialized(self): - try: - if self._ptr_array is not None: - return True - else: - return False - except AttributeError: - return False - - - - -@cython.cdivision(True) -@cython.boundscheck(False) -@cython.wraparound(False) -def _count_table_entries(ptr[:] count_array, ptr[:] ptr_array, - int64 element_size): - """ - For each bucket, count the number of elements in `ptr_array` and add - the count to the counts in `count_array`. - The element size gives the number of 32 bit elements per entry. - """ - cdef int64 length - cdef int64 count - cdef int64 bucket - cdef uint32* bucket_ptr - - for bucket in range(count_array.shape[0]): - bucket_ptr = (ptr_array[bucket]) - if bucket_ptr != NULL: - # First 64 bits are length of C-array - length = (bucket_ptr)[0] - count = (length - 2) // element_size - count_array[bucket] += count - - -@cython.boundscheck(False) -@cython.wraparound(False) -def _init_c_arrays(ptr[:] ptr_array, int64 element_size): - """ - Transform an array of counts into a pointer array, by replacing the - count in each element with a pointer to an initialized but empty - ``int32`` C-array. - The size of each C-array is the count mutliplied by the - `element_size`. - The first element of each C-array is is the currently filled size - of the C-array (an ``int64``) measured in number of ``int32`` - elements. - """ - cdef int64 bucket - cdef int64 count - cdef uint32* bucket_ptr - - for bucket in range(ptr_array.shape[0]): - # Before the C-array for a bucket initialized, the element in - # the pointer array contains the number of elements the C-array - # should hold - count = ptr_array[bucket] - if count != 0: - # Array size + n x element size - bucket_ptr = malloc( - (2 + count * element_size) * sizeof(uint32) - ) - if not bucket_ptr: - raise MemoryError() - # The initial size is 2, - # which is the size of the array size value (int64) - ( bucket_ptr)[0] = 2 - ptr_array[bucket] = bucket_ptr - - -@cython.boundscheck(False) -@cython.wraparound(False) -def _equal_c_arrays(ptr[:] self_ptr_array, ptr[:] other_ptr_array): - """ - Check if two pointer arrays are equal, i.e. they point to C-arrays - with equal elements. - """ - cdef int64 bucket - cdef int64 i - cdef int64 self_length, other_length - cdef uint32* self_bucket_ptr - cdef uint32* other_bucket_ptr - - for bucket in range(self_ptr_array.shape[0]): - self_bucket_ptr = self_ptr_array[bucket] - other_bucket_ptr = other_ptr_array[bucket] - if self_bucket_ptr != NULL or other_bucket_ptr != NULL: - if self_bucket_ptr == NULL or other_bucket_ptr == NULL: - # One of the tables has entries for this bucket - # while the other one has not - return False - # This bucket exists in both tables - self_length = (self_bucket_ptr )[0] - other_length = (other_bucket_ptr)[0] - if self_length != other_length: - return False - for i in range(2, self_length): - if self_bucket_ptr[i] != other_bucket_ptr[i]: - return False - - # If none of the previous checks failed, both objects are equal - return True - - -@cython.boundscheck(False) -@cython.wraparound(False) -def _append_entries(ptr[:] trg_ptr_array, ptr[:] src_ptr_array): - """ - Append the elements in all C-arrays of the source pointer array to - the corresponding C-arrays of the target pointer array. - - Expect that the target C-arrays are already initialized to - sufficient capacity. - """ - cdef int64 bucket - cdef int64 self_length, other_length, new_length - cdef uint32* self_kmer_ptr - cdef uint32* other_kmer_ptr - - for bucket in range(trg_ptr_array.shape[0]): - self_kmer_ptr = trg_ptr_array[bucket] - other_kmer_ptr = src_ptr_array[bucket] - if other_kmer_ptr != NULL: - self_length = (self_kmer_ptr)[0] - other_length = (other_kmer_ptr)[0] - # New new C-array needs the combucketed space of both - # arrays, but only one length value - new_length = self_length + other_length - 2 - (self_kmer_ptr)[0] = new_length - - # Append the entry from the other table - # to the entry in this table - self_kmer_ptr += self_length - other_kmer_ptr += 2 - memcpy( - self_kmer_ptr, other_kmer_ptr, - (other_length - 2) * sizeof(uint32) - ) - - -@cython.boundscheck(False) -@cython.wraparound(False) -def _pickle_c_arrays(ptr[:] ptr_array): - """ - Pickle the C arrays into a single concatenated :class:`ndarray`. - The lengths of each C-array on these concatenated array is saved as well. - """ - cdef int64 pointer_i, bucket_i, concat_i - cdef int64 length - cdef uint32* bucket_ptr - - # First pass: Count the total concatenated size - cdef int64 total_length = 0 - for pointer_i in range(ptr_array.shape[0]): - bucket_ptr = ptr_array[pointer_i] - if bucket_ptr != NULL: - # The first element of the C-array is the length - # of the array - total_length += (bucket_ptr)[0] - - # Second pass: Copy the C-arrays into a single concatenated array - # and track the start position of each C-array - cdef uint32[:] concatenated_array = np.empty(total_length, dtype=np.uint32) - cdef int64[:] lengths = np.empty(ptr_array.shape[0], dtype=np.int64) - concat_i = 0 - for pointer_i in range(ptr_array.shape[0]): - bucket_ptr = ptr_array[pointer_i] - if bucket_ptr != NULL: - length = (bucket_ptr)[0] - lengths[pointer_i] = length - memcpy( - &concatenated_array[concat_i], - bucket_ptr, - length * sizeof(uint32), - ) - concat_i += length - else: - lengths[pointer_i] = 0 - - return np.asarray(concatenated_array), np.asarray(lengths) - - -@cython.boundscheck(False) -@cython.wraparound(False) -def _unpickle_c_arrays(ptr[:] ptr_array, state): - """ - Unpickle the pickled `state` into the given `ptr_array`. - """ - cdef int64 pointer_i, concat_i - cdef int64 length - cdef uint32* bucket_ptr - - cdef uint32[:] concatenated_array = state[0] - cdef int64[:] lengths = state[1] - - concat_i = 0 - for pointer_i in range(ptr_array.shape[0]): - length = lengths[pointer_i] - if length != 0: - bucket_ptr = malloc(length * sizeof(uint32)) - if not bucket_ptr: - raise MemoryError - memcpy( - bucket_ptr, - &concatenated_array[concat_i], - length * sizeof(uint32), - ) - concat_i += length - ptr_array[pointer_i] = bucket_ptr - - -cdef inline void _deallocate_ptrs(ptr[:] ptrs): - cdef int64 kmer - for kmer in range(ptrs.shape[0]): - free(ptrs[kmer]) - - -cdef np.ndarray expand(np.ndarray array): - """ - Double the size of the first dimension of an existing 2D array. - """ - new_array = np.empty( - (array.shape[0] * 2, array.shape[1]), dtype=array.dtype - ) - new_array[:array.shape[0], :] = array - return new_array - - -def _prepare_mask(kmer_alphabet, ignore_mask, seq_length): - """ - Convert an ignore mask into a positive mask. - Multiple formats (boolean mask, pointer array, None) are supported - for the input. - """ - if ignore_mask is None: - kmer_mask = np.ones( - kmer_alphabet.kmer_array_length(seq_length), dtype=np.uint8 - ) - else: - if not isinstance(ignore_mask, np.ndarray): - raise TypeError( - f"The given mask is a '{type(ignore_mask).__name__}', " - f"but an ndarray was expected" - ) - if ignore_mask.dtype != np.dtype(bool): - raise ValueError("Expected a boolean mask") - if len(ignore_mask) != seq_length: - raise IndexError( - f"ignore mask has length {len(ignore_mask)}, " - f"but the length of the sequence is {seq_length}" - ) - kmer_mask = _to_kmer_mask( - np.frombuffer( - ignore_mask.astype(bool, copy=False), dtype=np.uint8 - ), - kmer_alphabet - ) - return kmer_mask - - -@cython.boundscheck(False) -@cython.wraparound(False) -def _to_kmer_mask(uint8[:] mask not None, kmer_alphabet): - """ - Transform a sequence ignore mask into a *k-mer* mask. - - The difference between those masks is that - - 1. the *k-mer* mask is shorter and - 2. a position *i* in the *k-mer* mask is false, if any - informative position of *k-mer[i]* is true in the ignore - mask. - """ - cdef int64 i, j - cdef bint is_retained - - cdef uint8[:] kmer_mask = np.empty( - kmer_alphabet.kmer_array_length(mask.shape[0]), dtype=np.uint8 - ) - cdef int64 offset - cdef int64 k = kmer_alphabet.k - cdef int64[:] spacing - - if kmer_alphabet.spacing is None: - # Continuous k-mers - for i in range(kmer_mask.shape[0]): - is_retained = True - # If any sequence position of this k-mer is removed, - # discard this k-mer position - for j in range(i, i + k): - if mask[j]: - is_retained = False - kmer_mask[i] = is_retained - - else: - # Spaced k-mers - spacing = kmer_alphabet.spacing - for i in range(kmer_mask.shape[0]): - is_retained = True - # If any sequence position of this k-mer is removed, - # discard this k-mer position - for j in range(spacing.shape[0]): - offset = spacing[j] - if mask[j + offset]: - is_retained = False - kmer_mask[i] = is_retained - - return np.asarray(kmer_mask) - - - -def _check_position_shape(position_arrays, kmer_arrays): - """ - Check if the given lists and each element have the same length - and raise an exception, if this is not the case. - """ - if len(position_arrays) != len(kmer_arrays): - raise IndexError( - f"{len(position_arrays)} position arrays " - f"for {len(kmer_arrays)} k-mer arrays were given" - ) - for i, (positions, kmers) in enumerate( - zip(position_arrays, kmer_arrays) - ): - if len(positions) != len(kmers): - raise IndexError( - f"{len(positions)} positions" - f"for {len(kmers)} k-mers were given at index {i}" - ) - - -def _check_same_kmer_alphabet(tables): - """ - Check if the *k-mer* alphabets of all tables are equal. - """ - ref_alph = tables[0].kmer_alphabet - for alph in (table.kmer_alphabet for table in tables): - if not alph == ref_alph: - raise ValueError( - "The *k-mer* alphabets of the tables are not equal " - "to each other" - ) - - -def _check_same_buckets(tables): - """ - Check if the bucket sizes of all tables are equal. - """ - ref_n_buckets = tables[0].n_buckets - for buckets in (table.n_buckets for table in tables): - if not buckets == ref_n_buckets: - raise ValueError( - "The number of buckets of the tables are not equal " - "to each other" - ) - - -def _check_kmer_bounds(kmers, kmer_alphabet): - """ - Check k-mer codes for out-of-bounds values. - """ - if np.any(kmers < 0) or np.any(kmers >= len(kmer_alphabet)): - raise AlphabetError( - "Given k-mer codes do not represent valid k-mers" - ) - - -def _check_multiple_kmer_bounds(kmer_arrays, kmer_alphabet): - """ - Check given arrays of k-mer codes for out-of-bounds values. - """ - for kmers in kmer_arrays: - if np.any(kmers < 0) or np.any(kmers >= len(kmer_alphabet)): - raise AlphabetError( - "Given k-mer codes do not represent valid k-mers" - ) - - -def _check_kmer_alphabet(kmer_alph): - """ - Check if the given object is a KmerAaphabet and raise an exception, - if this is not the case - """ - if not isinstance(kmer_alph, KmerAlphabet): - raise TypeError( - f"Got {type(kmer_alph).__name__}, but KmerAlphabet was expected" - ) - - -def _compute_masks(masks, sequences): - """ - Check, if the number of masks match the number of sequences, and - raise an exception if this is not the case. - If no masks are given, create a respective list of ``None`` values. - """ - if masks is None: - return [None] * len(sequences) - else: - if len(masks) != len(sequences): - raise IndexError( - f"{len(masks)} masks were given, " - f"but there are {len(sequences)} sequences" - ) - return masks - - -def _compute_ref_ids(ref_ids, sequences): - """ - Check, if the number of reference IDs match the number of - sequences, and raise an exception, if this is not the case. - If no reference IDs are given, create an array that simply - enumerates. - """ - if ref_ids is None: - return np.arange(len(sequences)) - else: - if len(ref_ids) != len(sequences): - raise IndexError( - f"{len(ref_ids)} reference IDs were given, " - f"but there are {len(sequences)} sequences" - ) - return ref_ids - - -def _compute_alphabet(given_alphabet, sequence_alphabets): - """ - If `given_alphabet` is None, find a common alphabet among - `sequence_alphabets` and raise an exception if this is not possible. - Otherwise just check compatibility of alphabets. - """ - if given_alphabet is None: - alphabet = common_alphabet(sequence_alphabets) - if alphabet is None: - raise ValueError( - "There is no common alphabet that extends all alphabets" - ) - return alphabet - else: - for alph in sequence_alphabets: - if not given_alphabet.extends(alph): - raise ValueError( - "The given alphabet is incompatible with a least one " - "alphabet of the given sequences" - ) - return given_alphabet - - -def _to_string(table): - lines = [] - for kmer in table.get_kmers(): - symbols = table.kmer_alphabet.decode(kmer) - if isinstance(table.alphabet, LetterAlphabet): - symbols = "".join(symbols) - else: - symbols = str(tuple(symbols)) - line = symbols + ": " + ", ".join( - [str((ref_id.item(), pos.item())) for ref_id, pos in table[kmer]] - ) - lines.append(line) - return "\n".join(lines) diff --git a/src/rust/lib.rs b/src/rust/lib.rs index 7e68d6189..14b631165 100644 --- a/src/rust/lib.rs +++ b/src/rust/lib.rs @@ -28,5 +28,6 @@ fn rust(module: &Bound<'_, PyModule>) -> PyResult<()> { &structure::module(module)?, "biotite.rust.structure", )?; + add_subpackage(module, &sequence::module(module)?, "biotite.rust.sequence")?; Ok(()) } diff --git a/src/rust/sequence/align/kmertable.rs b/src/rust/sequence/align/kmertable.rs new file mode 100644 index 000000000..4a9633a4f --- /dev/null +++ b/src/rust/sequence/align/kmertable.rs @@ -0,0 +1,2802 @@ +//! Containers for efficient retrieval of k-mer matches. +//! +//! `KmerTable` and `BucketKmerTable` map *k-mers* to the sequence positions where they appear. +//! Both share the same algorithms, which therefore live on a single generic core +//! (`GenericKmerTable`) parameterized by a `TableKind`. The kind selects +//! +//! 1. the entry type stored per position (`KmerEntry`) — a plain +//! `(ref_id, position)` for `KmerTable`, or a `(kmer, ref_id, position)` +//! for the bucketed variant, and +//! 2. how a *k-mer* maps to a table slot — the identity for `KmerTable` (one +//! slot per *k-mer*), or `kmer % n_buckets` for the bucketed variant. +//! +//! The backing store is a single `NestedArray`: one contiguous data buffer +//! plus an offsets array, so there is no per-*k-mer* allocation. + +use super::nested::{NestedArray, NestedArrayBuilder}; +use numpy::ndarray::Array2; +use numpy::{IntoPyArray, PyArray1, PyArray2, PyArrayMethods, PyReadonlyArray1, PyReadonlyArray2}; +use pyo3::exceptions::{PyIndexError, PyTypeError, PyValueError}; +use pyo3::prelude::*; +use pyo3::types::{PyBytes, PyDict, PyList, PyModule, PyString, PyTuple}; +use pyo3::Borrowed; +use std::mem::ManuallyDrop; + +// Label as a separate module to indicate that this exception comes +// from biotite +mod biotite { + pyo3::import_exception!(biotite.sequence, AlphabetError); +} + +/// A *k-mer* code. +pub type Kmer = i64; + +/// A thin wrapper around the Python `KmerAlphabet` class. +/// It re-exposes the Python methods to Rust. +struct KmerAlphabet { + wrapped: Py, +} + +impl KmerAlphabet { + /// Construct a new `KmerAlphabet` from a base `alphabet`, `k` and `spacing`. + fn new( + py: Python<'_>, + alphabet: Bound<'_, PyAny>, + k: usize, + spacing: Option>, + ) -> PyResult { + let module = PyModule::import(py, "biotite.sequence.align.kmeralphabet")?; + let kmer_alphabet = module + .getattr("KmerAlphabet")? + .call1((alphabet, k, spacing))?; + Ok(Self { + wrapped: kmer_alphabet.unbind(), + }) + } + + /// A clone of the wrapped Python object. + fn as_bound<'py>(&self, py: Python<'py>) -> Bound<'py, PyAny> { + self.wrapped.bind(py).clone() + } + + /// A new reference to the wrapped Python object. + fn clone_ref(&self, py: Python<'_>) -> Py { + self.wrapped.clone_ref(py) + } + + /// Whether this *k-mer* alphabet equals `other`. + fn equals(&self, py: Python<'_>, other: &KmerAlphabet) -> PyResult { + self.wrapped.bind(py).eq(other.wrapped.bind(py)) + } + + /// The number of possible *k-mers*. + fn len(&self, py: Python<'_>) -> PyResult { + self.wrapped.bind(py).len() + } + + /// The base alphabet the *k-mers* are composed of. + fn base_alphabet<'py>(&self, py: Python<'py>) -> PyResult> { + self.wrapped.bind(py).getattr("base_alphabet") + } + + /// Whether the base alphabet extends `alphabet`. + fn base_alphabet_extends(&self, py: Python<'_>, alphabet: &Bound<'_, PyAny>) -> PyResult { + self.base_alphabet(py)? + .call_method1("extends", (alphabet,))? + .extract() + } + + /// Decode a *k-mer* code into its symbols. + fn decode<'py>(&self, py: Python<'py>, kmer: Kmer) -> PyResult> { + self.wrapped.bind(py).call_method1("decode", (kmer,)) + } + + /// The overlapping *k-mer* codes of `sequence`. + fn create_kmers<'py>( + &self, + py: Python<'py>, + sequence: &Bound<'py, PyAny>, + ) -> PyResult> { + let sequence_code = sequence.getattr("code")?; + let py_array: Bound<'py, PyArray1> = self + .wrapped + .bind(py) + .call_method1("create_kmers", (sequence_code,))? + .extract()?; + Ok(py_array.readonly()) + } + + /// The number of overlapping *k-mers* in a sequence of `seq_length`. + fn kmer_array_length(&self, py: Python<'_>, seq_length: usize) -> PyResult { + self.wrapped + .bind(py) + .call_method1("kmer_array_length", (seq_length,))? + .extract() + } + + /// The length of the *k-mers*. + fn k(&self, py: Python<'_>) -> PyResult { + self.wrapped.bind(py).getattr("k")?.extract() + } + + /// The spacing model, if spaced *k-mers* are used. + fn spacing(&self, py: Python<'_>) -> PyResult>> { + let spacing = self.wrapped.bind(py).getattr("spacing")?; + if spacing.is_none() { + Ok(None) + } else { + let spacing_array: Bound<'_, PyArray1> = spacing.extract()?; + Ok(Some(spacing_array.to_vec()?)) + } + } +} + +/// A single *k-mer* match: one row of `N` `int64` columns of a match result. +struct Match([i64; N]); + +impl Match { + /// Convert collected matches into the expected `(n, N)` `int64` array. + fn into_array(matches: Vec, py: Python<'_>) -> Bound<'_, PyArray2> { + let n = matches.len(); + // `Vec>` is `Vec<[i64; N]>`, i.e. `n * N` contiguous `i64`. + // Reinterpret it as a flat `Vec` without copying: same allocation, + // same size and alignment. + let mut matches = ManuallyDrop::new(matches); + let flat = unsafe { + Vec::from_raw_parts( + matches.as_mut_ptr() as *mut i64, + n * N, + matches.capacity() * N, + ) + }; + Array2::from_shape_vec((n, N), flat) + .unwrap() + .into_pyarray(py) + } +} + +/// A position entry stored in a `GenericKmerTable`. +/// +/// Implementations carry exactly the data required by their `TableKind`: +/// `KmerTable` does not need to store the *k-mer* (it is the slot index), +/// while the bucketed variant does. +trait KmerEntry: Copy + Default + PartialEq { + /// Create an entry for the given `kmer` at `position` of sequence `ref_id`. + fn new(kmer: Kmer, ref_id: u32, position: u32) -> Self; + /// The reference ID of the indexed sequence. + fn ref_id(&self) -> u32; + /// The position within the indexed sequence. + fn position(&self) -> u32; + /// Whether this entry actually belongs to the given `kmer`. + /// + /// Always true for `KmerTable` (the slot already identifies the *k-mer*), + /// but the bucketed variant must compare the stored *k-mer*. + fn matches(&self, kmer: Kmer) -> bool; +} + +/// An entry of a `KmerTable`: just the position of a *k-mer*. +/// +/// The *k-mer* itself is not stored, as it is the slot index. `repr(C)` gives it +/// a stable layout for serialization (see `GenericKmerTable::get_state`). +#[derive(Copy, Clone, Default, PartialEq)] +#[repr(C)] +struct PlainEntry { + ref_id: u32, + position: u32, +} + +impl KmerEntry for PlainEntry { + #[inline(always)] + fn new(_kmer: Kmer, ref_id: u32, position: u32) -> Self { + PlainEntry { ref_id, position } + } + #[inline(always)] + fn ref_id(&self) -> u32 { + self.ref_id + } + #[inline(always)] + fn position(&self) -> u32 { + self.position + } + #[inline(always)] + fn matches(&self, _kmer: Kmer) -> bool { + true + } +} + +/// An entry of a `BucketKmerTable`: the *k-mer* alongside its position. +/// +/// As several distinct *k-mers* can share a bucket, the *k-mer* must be stored +/// to tell them apart on lookup. `repr(C)` gives it a stable layout for +/// serialization (see `GenericKmerTable::get_state`). +#[derive(Copy, Clone, Default, PartialEq)] +#[repr(C)] +struct BucketEntry { + kmer: Kmer, + ref_id: u32, + position: u32, +} + +impl KmerEntry for BucketEntry { + #[inline(always)] + fn new(kmer: Kmer, ref_id: u32, position: u32) -> Self { + BucketEntry { + kmer, + ref_id, + position, + } + } + #[inline(always)] + fn ref_id(&self) -> u32 { + self.ref_id + } + #[inline(always)] + fn position(&self) -> u32 { + self.position + } + #[inline(always)] + fn matches(&self, kmer: Kmer) -> bool { + self.kmer == kmer + } +} + +/// Selects the entry type and slot mapping of a `GenericKmerTable`. +trait TableKind: Sized { + /// The position entry stored per slot. + type Entry: KmerEntry; + + /// The slot index a `kmer` is stored in, given the number of slots. + fn slot_of(kmer: Kmer, n_slots: usize) -> usize; + + /// The *k-mer* represented by `entry` located in `slot`. + fn kmer_of_entry(slot: usize, entry: &Self::Entry) -> Kmer; + + /// The number of slots to allocate. + /// + /// `n_kmers` is the size of the *k-mer* alphabet, `total_kmers` the total + /// number of *k-mers* that will be stored (used to pick a default bucket + /// number) and `n_buckets` an explicitly requested bucket number. + fn resolve_n_slots( + py: Python<'_>, + n_kmers: usize, + total_kmers: usize, + n_buckets: Option, + ) -> PyResult; + + /// The number of entries in `slot` that belong to `kmer`. + /// + /// For `KmerTable` this is simply the slice length (the slot already + /// identifies the *k-mer*); the bucketed variant must filter by *k-mer*. + fn count_in_slot(slot: &[Self::Entry], kmer: Kmer) -> usize; + + /// All stored *k-mer* codes in ascending order. + fn collect_kmers(table: &GenericKmerTable) -> Vec; +} + +/// The `TableKind` of `KmerTable`: one slot per *k-mer*. +struct Plain; + +impl TableKind for Plain { + type Entry = PlainEntry; + + #[inline(always)] + fn slot_of(kmer: Kmer, _n_slots: usize) -> usize { + kmer as usize + } + + #[inline(always)] + fn kmer_of_entry(slot: usize, _entry: &PlainEntry) -> Kmer { + slot as Kmer + } + + fn resolve_n_slots( + _py: Python<'_>, + n_kmers: usize, + _total_kmers: usize, + _n_buckets: Option, + ) -> PyResult { + Ok(n_kmers) + } + + #[inline(always)] + fn count_in_slot(slot: &[PlainEntry], _kmer: Kmer) -> usize { + // Every entry in the slot belongs to the slot's k-mer + slot.len() + } + + fn collect_kmers(table: &GenericKmerTable) -> Vec { + // Each non-empty slot's index *is* a k-mer, and slots are ascending + (0..table.n_slots) + .filter(|&slot| !table.table[slot].is_empty()) + .map(|slot| slot as Kmer) + .collect() + } +} + +/// The `TableKind` of `BucketKmerTable`: *k-mers* are pooled into buckets. +struct Bucketed; + +impl TableKind for Bucketed { + type Entry = BucketEntry; + + #[inline(always)] + fn slot_of(kmer: Kmer, n_slots: usize) -> usize { + kmer as usize % n_slots + } + + #[inline(always)] + fn kmer_of_entry(_slot: usize, entry: &BucketEntry) -> Kmer { + entry.kmer + } + + fn resolve_n_slots( + py: Python<'_>, + n_kmers: usize, + total_kmers: usize, + n_buckets: Option, + ) -> PyResult { + let n_buckets = match n_buckets { + Some(n_buckets) => n_buckets, + None => { + // `bucket_number()` returns a NumPy float, so round it to an int + let n_buckets = PyModule::import(py, "biotite.sequence.align.buckets")? + .getattr("bucket_number")? + .call1((total_kmers,))?; + PyModule::import(py, "builtins")? + .getattr("int")? + .call1((n_buckets,))? + .extract()? + } + }; + // Never use more buckets than possible k-mers + Ok(n_buckets.min(n_kmers)) + } + + #[inline(always)] + fn count_in_slot(slot: &[BucketEntry], kmer: Kmer) -> usize { + slot.iter().filter(|entry| entry.matches(kmer)).count() + } + + fn collect_kmers(table: &GenericKmerTable) -> Vec { + // Distinct k-mers are scattered across buckets, so gather and sort them + let mut kmers: Vec = Vec::new(); + for slot in 0..table.n_slots { + for entry in &table.table[slot] { + kmers.push(entry.kmer); + } + } + kmers.sort_unstable(); + kmers.dedup(); + kmers + } +} + +impl<'a, 'py> FromPyObject<'a, 'py> for KmerAlphabet { + type Error = PyErr; + + fn extract(obj: Borrowed<'a, 'py, PyAny>) -> PyResult { + let kmer_alphabet_class = + PyModule::import(obj.py(), "biotite.sequence.align.kmeralphabet")? + .getattr("KmerAlphabet")?; + if !obj.is_instance(&kmer_alphabet_class)? { + return Err(PyTypeError::new_err(format!( + "Got {}, but KmerAlphabet was expected", + obj.get_type().name()? + ))); + } + Ok(KmerAlphabet { + wrapped: obj.to_owned().unbind(), + }) + } +} + +/// A source of `(kmer, ref_id, position)` triples for table construction. +/// +/// `KmerSupplier::for_each` is generic over the sink, so the +/// counting and filling passes in `GenericKmerTable::build` are each +/// monomorphized — there is no per-element dynamic dispatch. +trait KmerSupplier { + /// Feed every `(kmer, ref_id, position)` triple into `sink`. + /// + /// Called twice during construction and must yield identical triples both + /// times. `py` lets a supplier produce its triples lazily. + fn for_each(&self, py: Python<'_>, sink: F) -> PyResult<()>; +} + +/// Supplies a triple per *k-mer* of each input sequence (`from_sequences`). +/// +/// The *k-mers* are created lazily inside `for_each`, so only one *k-mer* array +/// is alive at any moment instead of one per sequence. +struct SequenceSupplier<'a, 'py> { + sequences: &'a [Bound<'py, PyAny>], + kmer_alphabet: KmerAlphabet, + ignore_masks: &'a [Option>], + ref_ids: &'a [u32], +} + +impl KmerSupplier for SequenceSupplier<'_, '_> { + fn for_each(&self, py: Python<'_>, mut sink: F) -> PyResult<()> { + for ((sequence, ignore_mask), &ref_id) in self + .sequences + .iter() + .zip(self.ignore_masks) + .zip(self.ref_ids) + { + let kmers = self.kmer_alphabet.create_kmers(py, sequence)?; + let kmers = kmers.as_slice()?; + match ignore_mask { + // No mask: every k-mer is included, no include mask is built + None => { + for (i, &kmer) in kmers.iter().enumerate() { + sink(kmer, ref_id, i as u32); + } + } + Some(_) => { + let mask = convert_ignore_into_include_mask( + py, + &self.kmer_alphabet, + ignore_mask.as_ref(), + sequence.len()?, + )?; + for (i, (&kmer, &included)) in kmers.iter().zip(&mask).enumerate() { + if included { + sink(kmer, ref_id, i as u32); + } + } + } + } + // The k-mer array is dropped here before the next sequence + } + Ok(()) + } +} + +/// Supplies a triple per given *k-mer*, using its index in the array as the +/// position, optionally filtered by a boolean include mask (`from_kmers`). +struct IndexedSupplier<'a> { + kmers: &'a [&'a [Kmer]], + masks: &'a [Vec], + ref_ids: &'a [u32], +} + +impl KmerSupplier for IndexedSupplier<'_> { + #[inline] + fn for_each(&self, _py: Python<'_>, mut sink: F) -> PyResult<()> { + for ((kmers, mask), &ref_id) in self.kmers.iter().zip(self.masks).zip(self.ref_ids) { + for (i, (&kmer, &included)) in kmers.iter().zip(mask).enumerate() { + if included { + sink(kmer, ref_id, i as u32); + } + } + } + Ok(()) + } +} + +/// Supplies triples with explicit positions and no mask (`from_kmer_selection`). +struct SelectedSupplier<'a> { + kmers: &'a [&'a [Kmer]], + positions: &'a [&'a [u32]], + ref_ids: &'a [u32], +} + +impl KmerSupplier for SelectedSupplier<'_> { + #[inline] + fn for_each(&self, _py: Python<'_>, mut sink: F) -> PyResult<()> { + for ((kmers, positions), &ref_id) in self.kmers.iter().zip(self.positions).zip(self.ref_ids) + { + for (&kmer, &position) in kmers.iter().zip(positions.iter()) { + sink(kmer, ref_id, position); + } + } + Ok(()) + } +} + +/// Supplies the triples from a `from_positions` mapping, iterating the original +/// `(ref_id, position)` arrays directly (no intermediate copy). +struct PositionSupplier<'a, 'py> { + entries: &'a [(Kmer, PyReadonlyArray2<'py, u32>)], +} + +impl KmerSupplier for PositionSupplier<'_, '_> { + #[inline] + fn for_each(&self, _py: Python<'_>, mut sink: F) -> PyResult<()> { + for (kmer, array) in self.entries { + for row in array.as_array().rows() { + sink(*kmer, row[0], row[1]); + } + } + Ok(()) + } +} + +/// Supplies the triples taken from the entries of existing tables (`from_tables`). +struct MergedSupplier<'a, K: TableKind> { + tables: &'a [&'a GenericKmerTable], +} + +impl KmerSupplier for MergedSupplier<'_, K> { + #[inline] + fn for_each(&self, _py: Python<'_>, mut sink: F) -> PyResult<()> { + for table in self.tables { + for slot in 0..table.n_slots { + for entry in &table.table[slot] { + sink( + K::kmer_of_entry(slot, entry), + entry.ref_id(), + entry.position(), + ); + } + } + } + Ok(()) + } +} + +/// The generic core shared by `KmerTable` and the bucketed variant. +/// +/// `K` selects the entry type and slot mapping (see `TableKind`). +struct GenericKmerTable { + /// The length of the *k-mers*. + k: usize, + /// The size of the *k-mer* alphabet (number of possible *k-mers*). + n_kmers: usize, + /// The number of table slots (equal to `n_kmers` for `KmerTable`). + n_slots: usize, + /// The wrapped Python `KmerAlphabet`. + kmer_alphabet: KmerAlphabet, + /// One inner array of position entries per slot. + table: NestedArray, +} + +impl GenericKmerTable { + /// Build a table by emitting every `(kmer, ref_id, position)` triple + /// twice: once to count occurrences per slot and once to scatter them into a + /// `NestedArrayBuilder` via counting sort. + fn build( + py: Python<'_>, + kmer_alphabet: KmerAlphabet, + k: usize, + n_kmers: usize, + n_slots: usize, + supplier: C, + ) -> PyResult { + // Count the number of entries per slot + // Reserve one extra element so `NestedArrayBuilder::new()` can append the + // final offset without reallocating + let mut counts: Vec = Vec::with_capacity(n_slots + 1); + counts.resize(n_slots, 0); + supplier.for_each(py, |kmer, _ref_id, _position| { + counts[K::slot_of(kmer, n_slots)] += 1; + })?; + + // Scatter the entries into their slots via counting sort + let mut builder = NestedArrayBuilder::new(counts); + supplier.for_each(py, |kmer, ref_id, position| { + // SAFETY: the count pass above counted exactly one entry per + // triple in slot `slot_of(kmer)`, and the same triples + // are replayed here, so no slot receives more entries than its + // length. + unsafe { + builder.push( + K::slot_of(kmer, n_slots), + ::new(kmer, ref_id, position), + ); + } + })?; + + Ok(Self { + k, + n_kmers, + n_slots, + kmer_alphabet, + table: builder.build(), + }) + } + + /// The entries stored for `kmer` (filtered by the *k-mer* for buckets). + #[inline] + fn lookup(&self, kmer: Kmer) -> impl Iterator { + self.table[K::slot_of(kmer, self.n_slots)] + .iter() + .filter(move |entry| entry.matches(kmer)) + } + + /// The reference IDs and positions stored for `kmer`, in storage order. + fn positions(&self, kmer: Kmer) -> Vec<(u32, u32)> { + self.lookup(kmer) + .map(|entry| (entry.ref_id(), entry.position())) + .collect() + } + + /// The number of stored positions for `kmer`. + #[inline] + fn count_kmer(&self, kmer: Kmer) -> usize { + K::count_in_slot(&self.table[K::slot_of(kmer, self.n_slots)], kmer) + } + + /// The *k-mer* codes similar to `kmer` according to a `SimilarityRule`. + fn similar_kmers( + &self, + py: Python<'_>, + rule: &Bound<'_, PyAny>, + kmer: Kmer, + ) -> PyResult> { + let array: Bound<'_, PyArray1> = rule + .call_method1("similar_kmers", (self.kmer_alphabet.as_bound(py), kmer))? + .extract()?; + Ok(array.readonly().as_slice()?.to_vec()) + } + + /// All *k-mer* codes with at least one stored position, in ascending order. + fn collect_kmers(&self) -> Vec { + K::collect_kmers(self) + } + + /// `from_sequences()` constructor, see the Python docstring. + #[allow(clippy::too_many_arguments)] + fn from_sequences( + py: Python<'_>, + k: usize, + sequences: Vec>, + ref_ids: Option>, + ignore_masks: Option>>>, + alphabet: Option>, + spacing: Option>, + n_buckets: Option, + ) -> PyResult { + let n_sequences = sequences.len(); + let ref_ids = compute_ref_ids(ref_ids, n_sequences)?; + let ignore_masks = compute_masks(ignore_masks, n_sequences)?; + let alphabet = get_alphabet(py, alphabet, &sequences)?; + let kmer_alphabet = KmerAlphabet::new(py, alphabet, k, spacing)?; + let n_kmers = kmer_alphabet.len(py)?; + + // The total (unmasked) number of k-mers, used to pick a bucket number. + // Computed without creating the k-mers, so no array is materialized here. + let total_kmers = sequences + .iter() + .map(|sequence| kmer_alphabet.kmer_array_length(py, sequence.len()?)) + .sum::>()?; + let n_slots = K::resolve_n_slots(py, n_kmers, total_kmers, n_buckets)?; + + // The k-mers are created lazily during construction, so only a single + // k-mer array is alive at a time. The supplier holds its own reference + // to the alphabet, leaving `kmer_alphabet` free to move into the table. + let supplier = SequenceSupplier { + sequences: &sequences, + kmer_alphabet: KmerAlphabet { + wrapped: kmer_alphabet.clone_ref(py), + }, + ignore_masks: &ignore_masks, + ref_ids: &ref_ids, + }; + Self::build(py, kmer_alphabet, k, n_kmers, n_slots, supplier) + } + + /// `from_kmers()` constructor, see the Python docstring. + fn from_kmers( + py: Python<'_>, + kmer_alphabet: KmerAlphabet, + kmers: Vec>, + ref_ids: Option>, + masks: Option>>>, + n_buckets: Option, + ) -> PyResult { + let k = kmer_alphabet.k(py)?; + let n_arrays = kmers.len(); + let ref_ids = compute_ref_ids(ref_ids, n_arrays)?; + let masks = compute_masks(masks, n_arrays)?; + + let n_kmers = kmer_alphabet.len(py)?; + let kmers_list: Vec> = kmers + .iter() + .map(|array| to_kmer_array(py, array)) + .collect::>()?; + let kmers_slices = as_slices(&kmers_list)?; + for kmers in &kmers_slices { + check_kmer_bounds(kmers, n_kmers)?; + } + + // Resolve each mask into an explicit boolean include mask + let masks_list: Vec> = masks + .iter() + .zip(kmers_slices.iter()) + .map(|(mask, kmers)| to_include_mask(mask.as_ref(), kmers.len())) + .collect::>()?; + + let total_kmers: usize = kmers_slices.iter().map(|s| s.len()).sum(); + let n_slots = K::resolve_n_slots(py, n_kmers, total_kmers, n_buckets)?; + + let supplier = IndexedSupplier { + kmers: &kmers_slices, + masks: &masks_list, + ref_ids: &ref_ids, + }; + Self::build(py, kmer_alphabet, k, n_kmers, n_slots, supplier) + } + + /// `from_kmer_selection()` constructor, see the Python docstring. + fn from_kmer_selection( + py: Python<'_>, + kmer_alphabet: KmerAlphabet, + positions: Vec>, + kmers: Vec>, + ref_ids: Option>, + n_buckets: Option, + ) -> PyResult { + let k = kmer_alphabet.k(py)?; + let n_kmers = kmer_alphabet.len(py)?; + + check_position_shape(&positions, &kmers)?; + let ref_ids = compute_ref_ids(ref_ids, kmers.len())?; + + let kmers_list: Vec> = kmers + .iter() + .map(|array| to_kmer_array(py, array)) + .collect::>()?; + let kmers_slices = as_slices(&kmers_list)?; + for kmers in &kmers_slices { + check_kmer_bounds(kmers, n_kmers)?; + } + let positions_list: Vec> = positions + .iter() + .map(|array| to_u32_array(py, array)) + .collect::>()?; + let positions_slices = as_slices(&positions_list)?; + + let total_kmers: usize = kmers_slices.iter().map(|s| s.len()).sum(); + let n_slots = K::resolve_n_slots(py, n_kmers, total_kmers, n_buckets)?; + + let supplier = SelectedSupplier { + kmers: &kmers_slices, + positions: &positions_slices, + ref_ids: &ref_ids, + }; + Self::build(py, kmer_alphabet, k, n_kmers, n_slots, supplier) + } + + /// `from_positions()` constructor, see the Python docstring. + fn from_positions( + py: Python<'_>, + kmer_alphabet: KmerAlphabet, + kmer_positions: &Bound<'_, PyDict>, + n_buckets: Option, + ) -> PyResult { + let k = kmer_alphabet.k(py)?; + let n_kmers = kmer_alphabet.len(py)?; + + // Collect and validate the position arrays, but iterate them directly + // during construction rather than materializing every triple + let mut entries: Vec<(Kmer, PyReadonlyArray2<'_, u32>)> = Vec::new(); + for (key, value) in kmer_positions.iter() { + let kmer: Kmer = key.extract()?; + if kmer < 0 || kmer as usize >= n_kmers { + return Err(biotite::AlphabetError::new_err(format!( + "k-mer code {kmer} does not represent a valid k-mer" + ))); + } + let array = to_u32_array2(py, &value)?; + let (n_rows, n_cols) = array.as_array().dim(); + if n_rows == 0 { + continue; + } + if n_cols != 2 { + return Err(PyIndexError::new_err(format!( + "Each entry in position array has {n_cols} values, but 2 were expected" + ))); + } + entries.push((kmer, array)); + } + + let total_kmers: usize = entries + .iter() + .map(|(_, array)| array.as_array().dim().0) + .sum(); + let n_slots = K::resolve_n_slots(py, n_kmers, total_kmers, n_buckets)?; + let supplier = PositionSupplier { entries: &entries }; + Self::build(py, kmer_alphabet, k, n_kmers, n_slots, supplier) + } + + /// `from_tables()` constructor, see the Python docstring. + fn from_tables(py: Python<'_>, tables: &[&GenericKmerTable]) -> PyResult { + if tables.is_empty() { + return Err(PyValueError::new_err("At least one table must be given")); + } + let first = tables[0]; + for table in &tables[1..] { + if !first.kmer_alphabet.equals(py, &table.kmer_alphabet)? { + return Err(PyValueError::new_err( + "The *k-mer* alphabets of the tables are not equal to each other", + )); + } + if first.n_slots != table.n_slots { + return Err(PyValueError::new_err( + "The number of buckets of the tables are not equal to each other", + )); + } + } + + let kmer_alphabet = KmerAlphabet { + wrapped: first.kmer_alphabet.clone_ref(py), + }; + let supplier = MergedSupplier { tables }; + Self::build( + py, + kmer_alphabet, + first.k, + first.n_kmers, + first.n_slots, + supplier, + ) + } + + /// `match_table()` method, see the Python docstring. + fn match_table<'py>( + &self, + py: Python<'py>, + other: &GenericKmerTable, + similarity_rule: Option<&Bound<'py, PyAny>>, + ) -> PyResult>> { + if !self.kmer_alphabet.equals(py, &other.kmer_alphabet)? { + return Err(PyValueError::new_err( + "The *k-mer* alphabets of the tables are not equal to each other", + )); + } + if self.n_slots != other.n_slots { + return Err(PyValueError::new_err( + "The number of buckets of the tables are not equal to each other", + )); + } + + let mut matches: Vec> = Vec::new(); + for slot in 0..other.n_slots { + for other_entry in &other.table[slot] { + let kmer = K::kmer_of_entry(slot, other_entry); + let other_ref = other_entry.ref_id() as i64; + let other_pos = other_entry.position() as i64; + if let Some(rule) = similarity_rule { + for sim_kmer in self.similar_kmers(py, rule, kmer)? { + for self_entry in self.lookup(sim_kmer) { + matches.push(Match([ + other_ref, + other_pos, + self_entry.ref_id() as i64, + self_entry.position() as i64, + ])); + } + } + } else { + for self_entry in self.lookup(kmer) { + matches.push(Match([ + other_ref, + other_pos, + self_entry.ref_id() as i64, + self_entry.position() as i64, + ])); + } + } + } + } + Ok(Match::into_array(matches, py)) + } + + /// `match()` method, see the Python docstring. + fn match_sequence<'py>( + &self, + py: Python<'py>, + sequence: &Bound<'py, PyAny>, + similarity_rule: Option<&Bound<'py, PyAny>>, + ignore_mask: Option<&Bound<'py, PyAny>>, + ) -> PyResult>> { + let seq_length = sequence.len()?; + if seq_length < self.k { + return Err(PyValueError::new_err("Sequence code is shorter than k")); + } + if !self + .kmer_alphabet + .base_alphabet_extends(py, &sequence.getattr("alphabet")?)? + { + return Err(PyValueError::new_err( + "The alphabet used for the k-mer index table is not equal to \ + the alphabet of the sequence", + )); + } + + let kmers = self.kmer_alphabet.create_kmers(py, sequence)?; + let kmers = kmers.as_slice()?; + let mask = + convert_ignore_into_include_mask(py, &self.kmer_alphabet, ignore_mask, seq_length)?; + + let mut matches: Vec> = Vec::new(); + for (i, &kmer) in kmers.iter().enumerate() { + if !mask[i] { + continue; + } + let query_pos = i as i64; + if let Some(rule) = similarity_rule { + for sim_kmer in self.similar_kmers(py, rule, kmer)? { + for entry in self.lookup(sim_kmer) { + matches.push(Match([ + query_pos, + entry.ref_id() as i64, + entry.position() as i64, + ])); + } + } + } else { + for entry in self.lookup(kmer) { + matches.push(Match([ + query_pos, + entry.ref_id() as i64, + entry.position() as i64, + ])); + } + } + } + Ok(Match::into_array(matches, py)) + } + + /// `match_kmer_selection()` method, see the Python docstring. + fn match_kmer_selection<'py>( + &self, + py: Python<'py>, + positions: &Bound<'py, PyAny>, + kmers: &Bound<'py, PyAny>, + ) -> PyResult>> { + let kmers = to_kmer_array(py, kmers)?; + let kmers = kmers.as_slice()?; + let positions = to_u32_array(py, positions)?; + let positions = positions.as_slice()?; + check_kmer_bounds(kmers, self.n_kmers)?; + if positions.len() != kmers.len() { + return Err(PyIndexError::new_err(format!( + "{} positions were given for {} k-mers", + positions.len(), + kmers.len() + ))); + } + + let mut matches: Vec> = Vec::new(); + for (&kmer, &seq_pos) in kmers.iter().zip(positions.iter()) { + for entry in self.lookup(kmer) { + matches.push(Match([ + seq_pos as i64, + entry.ref_id() as i64, + entry.position() as i64, + ])); + } + } + Ok(Match::into_array(matches, py)) + } + + /// `count()` for the given `kmers`. + fn count<'py>( + &self, + py: Python<'py>, + kmers: &Bound<'py, PyAny>, + ) -> PyResult>> { + let kmers = to_kmer_array(py, kmers)?; + let kmers = kmers.as_slice()?; + check_kmer_bounds(kmers, self.n_kmers)?; + let counts: Vec = kmers + .iter() + .map(|&kmer| self.count_kmer(kmer) as i64) + .collect(); + Ok(counts.into_pyarray(py)) + } + + /// `count()` for all *k-mers* in ascending order. + fn count_all<'py>(&self, py: Python<'py>) -> Bound<'py, PyArray1> { + let counts: Vec = (0..self.n_kmers as Kmer) + .map(|kmer| self.count_kmer(kmer) as i64) + .collect(); + counts.into_pyarray(py) + } + + /// `get_kmers()` method. + fn get_kmers<'py>(&self, py: Python<'py>) -> Bound<'py, PyArray1> { + self.collect_kmers().into_pyarray(py) + } + + /// `__getitem__()` method. + fn get_item<'py>(&self, py: Python<'py>, kmer: Kmer) -> PyResult>> { + if kmer < 0 || kmer as usize >= self.n_kmers { + return Err(biotite::AlphabetError::new_err(format!( + "k-mer code {kmer} is out of bounds for the given KmerAlphabet" + ))); + } + let positions = self.positions(kmer); + let mut flat: Vec = Vec::with_capacity(positions.len() * 2); + for (ref_id, position) in positions { + flat.push(ref_id); + flat.push(position); + } + let n = flat.len() / 2; + Ok(Array2::from_shape_vec((n, 2), flat) + .unwrap() + .into_pyarray(py)) + } + + /// `__contains__()` method. + fn contains(&self, kmer: Kmer) -> bool { + kmer >= 0 && (kmer as usize) < self.n_kmers && self.lookup(kmer).next().is_some() + } + + /// `__eq__()` method. + fn equals(&self, py: Python<'_>, other: &GenericKmerTable) -> PyResult { + // Equal k-mer alphabets imply equal `k` and (for `KmerTable`) `n_slots`; + // `n_slots` is still checked for the bucketed variant + if self.n_slots != other.n_slots { + return Ok(false); + } + if !self.kmer_alphabet.equals(py, &other.kmer_alphabet)? { + return Ok(false); + } + for slot in 0..self.n_slots { + if self.table[slot] != other.table[slot] { + return Ok(false); + } + } + Ok(true) + } + + /// `__str__()` method. + fn to_str(&self, py: Python<'_>) -> PyResult { + let base_alphabet = self.kmer_alphabet.base_alphabet(py)?; + let is_letter = is_letter_alphabet(py, &base_alphabet)?; + let builtins = PyModule::import(py, "builtins")?; + + let mut lines = Vec::new(); + for kmer in self.collect_kmers() { + let symbols = self.kmer_alphabet.decode(py, kmer)?; + let symbols = if is_letter { + PyString::new(py, "") + .call_method1("join", (symbols,))? + .extract::()? + } else { + let tuple = builtins.getattr("tuple")?.call1((symbols,))?; + builtins + .getattr("str")? + .call1((tuple,))? + .extract::()? + }; + let positions = self + .positions(kmer) + .iter() + .map(|(ref_id, position)| format!("({ref_id}, {position})")) + .collect::>() + .join(", "); + lines.push(format!("{symbols}: {positions}")); + } + Ok(lines.join("\n")) + } + + /// The picklable state: the offsets array (as `int64`) and the raw entry + /// buffer (see `from_state` for the inverse). + fn pickle_state<'py>( + &self, + py: Python<'py>, + ) -> (Bound<'py, PyArray1>, Bound<'py, PyBytes>) { + let (data, offsets) = self.table.raw_parts(); + let offsets: Vec = offsets.iter().map(|&o| o as i64).collect(); + // SAFETY: `data` is a live, initialized slice of `size_of_val(data)` + // bytes; the byte view borrows it and is only read (copied into the + // `PyBytes`) before `data` is released. `K::Entry` is `repr(C)` plain + // data, so its bytes are well-defined. + let bytes = unsafe { + std::slice::from_raw_parts(data.as_ptr() as *const u8, std::mem::size_of_val(data)) + }; + (offsets.into_pyarray(py), PyBytes::new(py, bytes)) + } + + /// Reconstruct a table from its pickled state (see `pickle_state`). + /// + /// This builds the table directly, so unpickling does not first allocate an + /// empty `n_slots`-sized table only to discard it. + fn from_state( + py: Python<'_>, + kmer_alphabet: KmerAlphabet, + offsets: &[i64], + bytes: &[u8], + ) -> PyResult { + if offsets.is_empty() { + return Err(PyValueError::new_err("Invalid state: empty offsets array")); + } + let k = kmer_alphabet.k(py)?; + let n_kmers = kmer_alphabet.len(py)?; + let n_slots = offsets.len() - 1; + let offsets: Vec = offsets.iter().map(|&o| o as usize).collect(); + + // Copy the entry bytes into an uninitialized buffer (no zero-init) + let entry_size = std::mem::size_of::(); + let total = bytes.len() / entry_size; + let mut data: Vec = Vec::with_capacity(total); + // SAFETY: All bit patterns are valid KmerEntry objects, + // so even from untrusted sources no invalid entries can be created + unsafe { + std::ptr::copy_nonoverlapping( + bytes.as_ptr(), + data.as_mut_ptr() as *mut u8, + total * entry_size, + ); + data.set_len(total); + } + + // SAFETY: in the worst case indexing the NestedArray will panic if the offsets are invalid + let table = unsafe { NestedArray::from_raw_parts(data, offsets) }; + Ok(Self { + k, + n_kmers, + n_slots, + kmer_alphabet, + table, + }) + } +} + +/// Whether the base alphabet is a `LetterAlphabet`. +fn is_letter_alphabet(py: Python<'_>, base_alphabet: &Bound<'_, PyAny>) -> PyResult { + let letter_alphabet = + PyModule::import(py, "biotite.sequence.alphabet")?.getattr("LetterAlphabet")?; + base_alphabet.is_instance(&letter_alphabet) +} + +/// Extract the contiguous slices of a list of read-only NumPy arrays. +fn as_slices<'a, T: numpy::Element>( + arrays: &'a [PyReadonlyArray1<'_, T>], +) -> PyResult> { + arrays + .iter() + .map(|array| array.as_slice().map_err(Into::into)) + .collect() +} + +/// Initialize the reference IDs, defaulting to `0..n` if not given. +fn compute_ref_ids(ref_ids: Option>, n: usize) -> PyResult> { + match ref_ids { + Some(ids) => { + if ids.len() != n { + return Err(PyIndexError::new_err(format!( + "{} reference IDs were given, but there are {} sequences", + ids.len(), + n + ))); + } + Ok(ids) + } + None => Ok((0..n as u32).collect()), + } +} + +/// Initialize the masks, defaulting to a list of `None` if not given. +fn compute_masks( + masks: Option>>>, + n: usize, +) -> PyResult>>> { + match masks { + Some(masks) => { + if masks.len() != n { + return Err(PyIndexError::new_err(format!( + "{} masks were given, but there are {} sequences", + masks.len(), + n + ))); + } + Ok(masks) + } + None => Ok((0..n).map(|_| None).collect()), + } +} + +/// Convert an ignore mask (over sequence positions) into a *k-mer* include mask. +/// +/// A *k-mer* is included if none of its informative positions is ignored. +fn convert_ignore_into_include_mask( + py: Python<'_>, + kmer_alphabet: &KmerAlphabet, + ignore_mask: Option<&Bound<'_, PyAny>>, + seq_length: usize, +) -> PyResult> { + let kmer_array_length = kmer_alphabet.kmer_array_length(py, seq_length)?; + let Some(mask_array) = ignore_mask else { + return Ok(vec![true; kmer_array_length]); + }; + + let py_array: Bound<'_, PyArray1> = mask_array.extract()?; + let mask = py_array.readonly(); + let mask = mask.as_slice()?; + if mask.len() != seq_length { + return Err(PyIndexError::new_err(format!( + "ignore mask has length {}, but the length of the sequence is {}", + mask.len(), + seq_length + ))); + } + + let k = kmer_alphabet.k(py)?; + let spacing = kmer_alphabet.spacing(py)?; + let mut include = vec![true; kmer_array_length]; + match spacing { + None => { + // Continuous k-mers cover positions i..i+k + for (i, included) in include.iter_mut().enumerate() { + if mask[i..i + k].iter().any(|&ignored| ignored) { + *included = false; + } + } + } + Some(spacing) => { + // Spaced k-mers cover positions i + offset for each model offset + for (i, included) in include.iter_mut().enumerate() { + if spacing.iter().any(|&offset| mask[i + offset as usize]) { + *included = false; + } + } + } + } + Ok(include) +} + +/// Resolve a single mask into an explicit boolean include mask of `length`. +fn to_include_mask(mask: Option<&Bound<'_, PyAny>>, length: usize) -> PyResult> { + match mask { + None => Ok(vec![true; length]), + Some(mask) => { + let array: Bound<'_, PyArray1> = mask.extract()?; + Ok(array.readonly().as_slice()?.to_vec()) + } + } +} + +/// Find or check the common alphabet of the input `sequences`. +fn get_alphabet<'py>( + py: Python<'py>, + given_alphabet: Option>, + sequences: &[Bound<'py, PyAny>], +) -> PyResult> { + let module = PyModule::import(py, "biotite.sequence.alphabet")?; + let sequence_alphabets: Vec> = sequences + .iter() + .map(|seq| seq.getattr("alphabet")) + .collect::>()?; + + match given_alphabet { + Some(alphabet) => { + for sequence_alphabet in &sequence_alphabets { + if !alphabet + .call_method1("extends", (sequence_alphabet,))? + .extract::()? + { + return Err(PyValueError::new_err( + "The given alphabet is incompatible with a least one \ + alphabet of the given sequences", + )); + } + } + Ok(alphabet) + } + None => { + let alphabet = module + .getattr("common_alphabet")? + .call1((sequence_alphabets,))?; + if alphabet.is_none() { + return Err(PyValueError::new_err( + "There is no common alphabet that extends all alphabets", + )); + } + Ok(alphabet) + } + } +} + +/// Coerce a Python object into a contiguous `int64` k-mer array (no copy if +/// already `int64`). +fn to_kmer_array<'py>( + py: Python<'py>, + array: &Bound<'py, PyAny>, +) -> PyResult> { + let np = PyModule::import(py, "numpy")?; + let array: Bound<'py, PyArray1> = np + .call_method1("ascontiguousarray", (array, np.getattr("int64")?))? + .extract()?; + Ok(array.readonly()) +} + +/// Coerce a Python object into a contiguous `uint32` array (no copy if already +/// `uint32`). +fn to_u32_array<'py>( + py: Python<'py>, + array: &Bound<'py, PyAny>, +) -> PyResult> { + let np = PyModule::import(py, "numpy")?; + let array: Bound<'py, PyArray1> = np + .call_method1("ascontiguousarray", (array, np.getattr("uint32")?))? + .extract()?; + Ok(array.readonly()) +} + +/// Coerce a Python object into a contiguous `uint32` 2D array. +fn to_u32_array2<'py>( + py: Python<'py>, + array: &Bound<'py, PyAny>, +) -> PyResult> { + let np = PyModule::import(py, "numpy")?; + let array: Bound<'py, PyArray2> = np + .call_method1("ascontiguousarray", (array, np.getattr("uint32")?))? + .extract()?; + Ok(array.readonly()) +} + +/// Raise an `AlphabetError` if any *k-mer* code is out of bounds. +fn check_kmer_bounds(kmers: &[Kmer], n_kmers: usize) -> PyResult<()> { + // Out-of-bounds k-mers are not expected, so mark the error path `#[cold]`: + // this steers the branch predictor and code layout towards the common + // in-bounds case. + #[cold] + #[inline(never)] + fn out_of_bounds() -> PyErr { + biotite::AlphabetError::new_err("Given k-mer codes do not represent valid k-mers") + } + for &kmer in kmers { + if kmer < 0 || kmer as usize >= n_kmers { + return Err(out_of_bounds()); + } + } + Ok(()) +} + +/// Check that the number of position arrays and k-mer arrays match elementwise. +fn check_position_shape( + positions: &[Bound<'_, PyAny>], + kmers: &[Bound<'_, PyAny>], +) -> PyResult<()> { + if positions.len() != kmers.len() { + return Err(PyIndexError::new_err(format!( + "{} position arrays for {} k-mer arrays were given", + positions.len(), + kmers.len() + ))); + } + for (i, (positions, kmers)) in positions.iter().zip(kmers.iter()).enumerate() { + if positions.len()? != kmers.len()? { + return Err(PyIndexError::new_err(format!( + "{} positions for {} k-mers were given at index {}", + positions.len()?, + kmers.len()?, + i + ))); + } + } + Ok(()) +} + +/// This class represents a *k-mer* index table. +/// It maps *k-mers* (subsequences with length *k*) to the sequence +/// positions, where the *k-mer* appears. +/// It is primarily used to find *k-mer* matches between two sequences. +/// A match is defined as a *k-mer* that appears in both sequences. +/// Instances of this class are immutable. +/// +/// There are multiple ways to create a :class:`KmerTable`: +/// +/// - :meth:`from_sequences()` iterates through all overlapping +/// *k-mers* in a sequence and stores the sequence position of +/// each *kmer* in the table. +/// - :meth:`from_kmers()` is similar to :meth:`from_sequences()` +/// but directly accepts *k-mers* as input instead of sequences. +/// - :meth:`from_kmer_selection()` takes a combination of *k-mers* +/// and their positions in a sequence, which can be used to +/// apply subset selectors, such as :class:`MinimizerSelector`. +/// - :meth:`from_tables()` merges the entries from multiple +/// :class:`KmerTable` objects into a new table. +/// - :meth:`from_positions()` let's the user provide manual +/// *k-mer* positions, which can be useful for loading a +/// :class:`KmerTable` from file. +/// +/// Each indexed *k-mer* position is represented by a tuple of +/// +/// 1. a unique reference ID that identifies to which sequence a +/// position refers to and +/// 2. the zero-based sequence position of the first symbol in the +/// *k-mer*. +/// +/// The :meth:`match()` method iterates through all overlapping *k-mers* +/// in another sequence and, for each *k-mer*, looks up the reference +/// IDs and positions of this *k-mer* in the table. +/// For each matching position, it adds the *k-mer* position in this +/// sequence, the matching reference ID and the matching sequence +/// position to the array of matches. +/// Finally these matches are returned to the user. +/// Optionally, a :class:`SimilarityRule` can be supplied, to find +/// also matches for similar *k-mers*. +/// This is especially useful for protein sequences to match two +/// *k-mers* with a high substitution probability. +/// +/// The positions for a given *k-mer* code can be obtained via indexing. +/// Iteration over a :class:`KmerTable` yields the *k-mers* that have at +/// least one associated position. +/// The *k-mer* code for a *k-mer* can be calculated with +/// ``table.kmer_alphabet.encode()`` (see :class:`KmerAlphabet`). +/// +/// Attributes +/// ---------- +/// kmer_alphabet : KmerAlphabet +/// The internal :class:`KmerAlphabet`, that is used to +/// encode all overlapping *k-mers* of an input sequence. +/// alphabet : Alphabet +/// The base alphabet, from which this :class:`KmerTable` was +/// created. +/// k : int +/// The length of the *k-mers*. +/// +/// See Also +/// -------- +/// BucketKmerTable +/// +/// Notes +/// ----- +/// +/// The design of the :class:`KmerTable` is inspired by the *MMseqs2* +/// software :footcite:`Steinegger2017`. +/// +/// *Memory consumption* +/// +/// For efficient mapping, a :class:`KmerTable` contains two large arrays: +/// +/// 1. An array that holds all *k-mer* positions. +/// 2. An accompanying array that maps each *k-mer* to the index range in the +/// first array, where the corresponding *k-mer* positions are stored. +/// +/// The required memory space :math:`S` in byte is +/// +/// .. math:: +/// +/// S = 8 n^k + 8L, +/// +/// where :math:`n` is the number of symbols in the alphabet and +/// :math:`L` is the summed length of all sequences added to the table. +/// +/// *Multiprocessing* +/// +/// :class:`KmerTable` objects can be used in multi-processed setups: +/// Adding a large database of sequences to a table can be sped up by +/// splitting the database into smaller chunks and create a separate +/// table for each chunk in separate processes. +/// Eventually, the tables can be merged to one large table using +/// :meth:`from_tables()`. +/// +/// Since :class:`KmerTable` supports the *pickle* protocol, +/// the matching step can also be divided into multiple processes, if +/// multiple sequences need to be matched. +/// +/// *Storage on hard drive* +/// +/// The most time efficient way to read/write a :class:`KmerTable` is +/// the *pickle* format. +/// If a custom format is desired, the user needs to extract the +/// reference IDs and position for each *k-mer*. +/// To restrict this task to all *k-mers* that have at least one match +/// :meth:`get_kmers()` can be used. +/// Conversely, the reference IDs and positions can be restored via +/// :meth:`from_positions()`. +/// +/// References +/// ---------- +/// +/// .. footbibliography:: +/// +/// Examples +/// -------- +/// +/// Create a *2-mer* index table for some nucleotide sequences: +/// +/// >>> table = KmerTable.from_sequences( +/// ... k = 2, +/// ... sequences = [NucleotideSequence("TTATA"), NucleotideSequence("CTAG")], +/// ... ref_ids = [0, 1] +/// ... ) +/// +/// Display the contents of the table as +/// (reference ID, sequence position) tuples: +/// +/// >>> print(table) +/// AG: (1, 2) +/// AT: (0, 2) +/// CT: (1, 0) +/// TA: (0, 1), (0, 3), (1, 1) +/// TT: (0, 0) +/// +/// Find matches of the table with a sequence: +/// +/// >>> query = NucleotideSequence("TAG") +/// >>> matches = table.match(query) +/// >>> for query_pos, table_ref_id, table_pos in matches: +/// ... print("Query sequence position:", query_pos) +/// ... print("Table reference ID: ", table_ref_id) +/// ... print("Table sequence position:", table_pos) +/// ... print() +/// Query sequence position: 0 +/// Table reference ID: 0 +/// Table sequence position: 1 +/// +/// Query sequence position: 0 +/// Table reference ID: 0 +/// Table sequence position: 3 +/// +/// Query sequence position: 0 +/// Table reference ID: 1 +/// Table sequence position: 1 +/// +/// Query sequence position: 1 +/// Table reference ID: 1 +/// Table sequence position: 2 +/// +/// +/// Get all reference IDs and positions for a given *k-mer*: +/// +/// >>> kmer_code = table.kmer_alphabet.encode("TA") +/// >>> print(table[kmer_code]) +/// [[0 1] +/// [0 3] +/// [1 1]] +#[pyclass(module = "biotite.rust.sequence.align")] +pub struct KmerTable(GenericKmerTable); + +#[pymethods] +impl KmerTable { + /// Reconstruct a table from its pickled state (called by `__reduce__`). + #[staticmethod] + fn _from_state( + py: Python<'_>, + kmer_alphabet: KmerAlphabet, + offsets: PyReadonlyArray1<'_, i64>, + data: Bound<'_, PyBytes>, + ) -> PyResult { + Ok(KmerTable(GenericKmerTable::from_state( + py, + kmer_alphabet, + offsets.as_slice()?, + data.as_bytes(), + )?)) + } + + #[getter] + fn kmer_alphabet(&self, py: Python<'_>) -> Py { + self.0.kmer_alphabet.clone_ref(py) + } + + #[getter] + fn alphabet<'py>(&self, py: Python<'py>) -> PyResult> { + self.0.kmer_alphabet.base_alphabet(py) + } + + #[getter] + fn k(&self) -> usize { + self.0.k + } + + /// Create a :class:`KmerTable` by storing the positions of all + /// overlapping *k-mers* from the input `sequences`. + /// + /// Parameters + /// ---------- + /// k : int + /// The length of the *k-mers*. + /// sequences : sized iterable object of Sequence, length=m + /// The sequences to get the *k-mer* positions from. + /// These sequences must have equal alphabets, or one of these + /// sequences must have an alphabet that extends the alphabets + /// of all other sequences. + /// ref_ids : sized iterable object of int, length=m, optional + /// The reference IDs for the given sequences. + /// These are used to identify the corresponding sequence for a + /// *k-mer* match. + /// By default the IDs are counted from *0* to *m*. + /// ignore_masks : sized iterable object of (ndarray, dtype=bool), length=m, optional + /// Sequence positions to ignore. + /// *k-mers* that involve these sequence positions are not added + /// to the table. + /// This is used e.g. to skip repeat regions. + /// If provided, the list must contain one boolean mask + /// (or ``None``) for each sequence, and each bolean mask must + /// have the same length as the sequence. + /// By default, no sequence position is ignored. + /// alphabet : Alphabet, optional + /// The alphabet to use for this table. + /// It must extend the alphabets of the input `sequences`. + /// By default, an appropriate alphabet is inferred from the + /// input `sequences`. + /// This option is usually used for compatibility with another + /// sequence/table in the matching step. + /// spacing : None or str or list or ndarray, dtype=int, shape=(k,) + /// If provided, spaced *k-mers* are used instead of continuous + /// ones. + /// The value contains the *informative* positions relative to + /// the start of the *k-mer*, also called the *model*. + /// The number of *informative* positions must equal *k*. + /// Refer to :class:`KmerAlphabet` for more details. + /// + /// See Also + /// -------- + /// from_kmers : The same functionality based on already created *k-mers* + /// + /// Returns + /// ------- + /// table : KmerTable + /// The newly created table. + /// + /// Examples + /// -------- + /// + /// >>> sequences = [NucleotideSequence("TTATA"), NucleotideSequence("CTAG")] + /// >>> table = KmerTable.from_sequences( + /// ... 2, sequences, ref_ids=[100, 101] + /// ... ) + /// >>> print(table) + /// AG: (101, 2) + /// AT: (100, 2) + /// CT: (101, 0) + /// TA: (100, 1), (100, 3), (101, 1) + /// TT: (100, 0) + /// + /// Give an explicit compatible alphabet: + /// + /// >>> table = KmerTable.from_sequences( + /// ... 2, sequences, ref_ids=[100, 101], + /// ... alphabet=NucleotideSequence.ambiguous_alphabet() + /// ... ) + /// + /// Ignore all ``N`` in a sequence: + /// + /// >>> sequence = NucleotideSequence("ACCNTANNG") + /// >>> table = KmerTable.from_sequences( + /// ... 2, [sequence], ignore_masks=[sequence.symbols == "N"] + /// ... ) + /// >>> print(table) + /// AC: (0, 0) + /// CC: (0, 1) + /// TA: (0, 4) + #[staticmethod] + #[pyo3(signature = (k, sequences, ref_ids=None, ignore_masks=None, alphabet=None, spacing=None))] + fn from_sequences( + py: Python<'_>, + k: usize, + sequences: Vec>, + ref_ids: Option>, + ignore_masks: Option>>>, + alphabet: Option>, + spacing: Option>, + ) -> PyResult { + Ok(KmerTable(GenericKmerTable::from_sequences( + py, + k, + sequences, + ref_ids, + ignore_masks, + alphabet, + spacing, + None, + )?)) + } + + /// Create a :class:`KmerTable` by storing the positions of all + /// input *k-mers*. + /// + /// Parameters + /// ---------- + /// kmer_alphabet : KmerAlphabet + /// The :class:`KmerAlphabet` to use for the new table. + /// Should be the same alphabet that was used to calculate the + /// input *kmers*. + /// kmers : sized iterable object of (ndarray, dtype=np.int64), length=m + /// List where each array contains the *k-mer* codes from a + /// sequence. + /// For each array the index of the *k-mer* code in the array + /// is stored in the table as sequence position. + /// ref_ids : sized iterable object of int, length=m, optional + /// The reference IDs for the sequences. + /// These are used to identify the corresponding sequence for a + /// *k-mer* match. + /// By default the IDs are counted from *0* to *m*. + /// masks : sized iterable object of (ndarray, dtype=bool), length=m, optional + /// A *k-mer* code at a position, where the corresponding mask + /// is false, is not added to the table. + /// By default, all positions are added. + /// + /// See Also + /// -------- + /// from_sequences : The same functionality based on undecomposed sequences + /// + /// Returns + /// ------- + /// table : KmerTable + /// The newly created table. + /// + /// Examples + /// -------- + /// + /// >>> sequences = [ProteinSequence("BIQTITE"), ProteinSequence("NIQBITE")] + /// >>> kmer_alphabet = KmerAlphabet(ProteinSequence.alphabet, 3) + /// >>> kmer_codes = [kmer_alphabet.create_kmers(s.code) for s in sequences] + /// >>> for code in kmer_codes: + /// ... print(code) + /// [11701 4360 7879 9400 4419] + /// [ 6517 4364 7975 11704 4419] + /// >>> table = KmerTable.from_kmers( + /// ... kmer_alphabet, kmer_codes + /// ... ) + /// >>> print(table) + /// IQT: (0, 1) + /// IQB: (1, 1) + /// ITE: (0, 4), (1, 4) + /// NIQ: (1, 0) + /// QTI: (0, 2) + /// QBI: (1, 2) + /// TIT: (0, 3) + /// BIQ: (0, 0) + /// BIT: (1, 3) + #[staticmethod] + #[pyo3(signature = (kmer_alphabet, kmers, ref_ids=None, masks=None))] + fn from_kmers( + py: Python<'_>, + kmer_alphabet: KmerAlphabet, + kmers: Vec>, + ref_ids: Option>, + masks: Option>>>, + ) -> PyResult { + Ok(KmerTable(GenericKmerTable::from_kmers( + py, + kmer_alphabet, + kmers, + ref_ids, + masks, + None, + )?)) + } + + /// Create a :class:`KmerTable` by storing the positions of a + /// filtered subset of input *k-mers*. + /// + /// This can be used to reduce the number of stored *k-mers* using + /// a *k-mer* subset selector such as :class:`MinimizerSelector`. + /// + /// Parameters + /// ---------- + /// kmer_alphabet : KmerAlphabet + /// The :class:`KmerAlphabet` to use for the new table. + /// Should be the same alphabet that was used to calculate the + /// input *kmers*. + /// positions : sized iterable object of (ndarray, shape=(n,), dtype=uint32), length=m + /// List where each array contains the sequence positions of + /// the filtered subset of *k-mers* given in `kmers`. + /// The list may contain multiple elements for multiple + /// sequences. + /// kmers : sized iterable object of (ndarray, shape=(n,), dtype=np.int64), length=m + /// List where each array contains the filtered subset of + /// *k-mer* codes from a sequence. + /// For each array the index of the *k-mer* code in the array, + /// is stored in the table as sequence position. + /// The list may contain multiple elements for multiple + /// sequences. + /// ref_ids : sized iterable object of int, length=m, optional + /// The reference IDs for the sequences. + /// These are used to identify the corresponding sequence for a + /// *k-mer* match. + /// By default the IDs are counted from *0* to *m*. + /// + /// Returns + /// ------- + /// table : KmerTable + /// The newly created table. + /// + /// Examples + /// -------- + /// + /// Reduce the size of sequence data in the table using minimizers: + /// + /// >>> sequence1 = ProteinSequence("THIS*IS*A*SEQVENCE") + /// >>> kmer_alph = KmerAlphabet(sequence1.alphabet, k=3) + /// >>> minimizer = MinimizerSelector(kmer_alph, window=4) + /// >>> minimizer_pos, minimizers = minimizer.select(sequence1) + /// >>> kmer_table = KmerTable.from_kmer_selection( + /// ... kmer_alph, [minimizer_pos], [minimizers] + /// ... ) + /// + /// Use the same :class:`MinimizerSelector` to select the minimizers + /// from the query sequence and match them against the table. + /// Although the amount of *k-mers* is reduced, matching is still + /// guanrateed to work, if the two sequences share identity in the + /// given window: + /// + /// >>> sequence2 = ProteinSequence("ANQTHER*SEQVENCE") + /// >>> minimizer_pos, minimizers = minimizer.select(sequence2) + /// >>> matches = kmer_table.match_kmer_selection(minimizer_pos, minimizers) + /// >>> print(matches) + /// [[ 9 0 11] + /// [12 0 14]] + /// >>> for query_pos, _, db_pos in matches: + /// ... print(sequence1) + /// ... print(" " * (db_pos-1) + "^" * kmer_table.k) + /// ... print(sequence2) + /// ... print(" " * (query_pos-1) + "^" * kmer_table.k) + /// ... print() + /// THIS*IS*A*SEQVENCE + /// ^^^ + /// ANQTHER*SEQVENCE + /// ^^^ + /// + /// THIS*IS*A*SEQVENCE + /// ^^^ + /// ANQTHER*SEQVENCE + /// ^^^ + /// + #[staticmethod] + #[pyo3(signature = (kmer_alphabet, positions, kmers, ref_ids=None))] + fn from_kmer_selection( + py: Python<'_>, + kmer_alphabet: KmerAlphabet, + positions: Vec>, + kmers: Vec>, + ref_ids: Option>, + ) -> PyResult { + Ok(KmerTable(GenericKmerTable::from_kmer_selection( + py, + kmer_alphabet, + positions, + kmers, + ref_ids, + None, + )?)) + } + + /// Create a :class:`KmerTable` by merging the *k-mer* positions + /// from existing `tables`. + /// + /// Parameters + /// ---------- + /// tables : iterable object of KmerTable + /// The tables to be merged. + /// All tables must have equal :class:`KmerAlphabet` objects, + /// i.e. the same *k* and equal base alphabets. + /// + /// Returns + /// ------- + /// table : KmerTable + /// The newly created table. + /// + /// Examples + /// -------- + /// + /// >>> table1 = KmerTable.from_sequences( + /// ... 2, [NucleotideSequence("TTATA")], ref_ids=[100] + /// ... ) + /// >>> table2 = KmerTable.from_sequences( + /// ... 2, [NucleotideSequence("CTAG")], ref_ids=[101] + /// ... ) + /// >>> merged_table = KmerTable.from_tables([table1, table2]) + /// >>> print(merged_table) + /// AG: (101, 2) + /// AT: (100, 2) + /// CT: (101, 0) + /// TA: (100, 1), (100, 3), (101, 1) + /// TT: (100, 0) + #[staticmethod] + fn from_tables(py: Python<'_>, tables: Vec>) -> PyResult { + let refs: Vec<&GenericKmerTable> = tables.iter().map(|table| &table.0).collect(); + Ok(KmerTable(GenericKmerTable::from_tables(py, &refs)?)) + } + + /// Create a :class:`KmerTable` from *k-mer* reference IDs and + /// positions. + /// This constructor is especially useful for restoring a table + /// from previously serialized data. + /// + /// Parameters + /// ---------- + /// kmer_alphabet : KmerAlphabet + /// The :class:`KmerAlphabet` to use for the new table + /// kmer_positions : dict of (int -> ndarray, shape=(n,2), dtype=int) + /// A dictionary representing the *k-mer* reference IDs and + /// positions to be stored in the newly created table. + /// It maps a *k-mer* code to a :class:`ndarray`. + /// To achieve a high performance the data type ``uint32`` + /// is preferred for the arrays. + /// + /// Returns + /// ------- + /// table : KmerTable + /// The newly created table. + /// + /// Examples + /// -------- + /// + /// >>> sequence = ProteinSequence("BIQTITE") + /// >>> table = KmerTable.from_sequences(3, [sequence], ref_ids=[100]) + /// >>> print(table) + /// IQT: (100, 1) + /// ITE: (100, 4) + /// QTI: (100, 2) + /// TIT: (100, 3) + /// BIQ: (100, 0) + /// >>> data = {kmer: table[kmer] for kmer in table} + /// >>> print(data) + /// {4360: array([[100, 1]], dtype=uint32), 4419: array([[100, 4]], dtype=uint32), 7879: array([[100, 2]], dtype=uint32), 9400: array([[100, 3]], dtype=uint32), 11701: array([[100, 0]], dtype=uint32)} + /// >>> restored_table = KmerTable.from_positions(table.kmer_alphabet, data) + /// >>> print(restored_table) + /// IQT: (100, 1) + /// ITE: (100, 4) + /// QTI: (100, 2) + /// TIT: (100, 3) + /// BIQ: (100, 0) + #[staticmethod] + fn from_positions( + py: Python<'_>, + kmer_alphabet: KmerAlphabet, + kmer_positions: Bound<'_, PyDict>, + ) -> PyResult { + Ok(KmerTable(GenericKmerTable::from_positions( + py, + kmer_alphabet, + &kmer_positions, + None, + )?)) + } + + /// Find matches between the *k-mers* in this table with the + /// *k-mers* in another `table`. + /// + /// This means that for each *k-mer* the cartesian product between + /// the positions in both tables is added to the matches. + /// + /// Parameters + /// ---------- + /// table : KmerTable + /// The table to be matched. + /// Both tables must have equal :class:`KmerAlphabet` objects, + /// i.e. the same *k* and equal base alphabets. + /// similarity_rule : SimilarityRule, optional + /// If this parameter is given, not only exact *k-mer* matches + /// are considered, but also similar ones according to the given + /// :class:`SimilarityRule`. + /// + /// Returns + /// ------- + /// matches : ndarray, shape=(n,4), dtype=np.int64 + /// The *k-mer* matches. + /// Each row contains one match. Each match has the following + /// columns: + /// + /// 0. The reference ID of the matched sequence in the other + /// table + /// 1. The sequence position of the matched sequence in the + /// other table + /// 2. The reference ID of the matched sequence in this + /// table + /// 3. The sequence position of the matched sequence in this + /// table + /// + /// Notes + /// ----- + /// + /// There is no guaranteed order of the reference IDs or + /// sequence positions in the returned matches. + /// + /// Examples + /// -------- + /// + /// >>> sequence1 = ProteinSequence("BIQTITE") + /// >>> table1 = KmerTable.from_sequences(3, [sequence1], ref_ids=[100]) + /// >>> print(table1) + /// IQT: (100, 1) + /// ITE: (100, 4) + /// QTI: (100, 2) + /// TIT: (100, 3) + /// BIQ: (100, 0) + /// >>> sequence2 = ProteinSequence("TITANITE") + /// >>> table2 = KmerTable.from_sequences(3, [sequence2], ref_ids=[101]) + /// >>> print(table2) + /// ANI: (101, 3) + /// ITA: (101, 1) + /// ITE: (101, 5) + /// NIT: (101, 4) + /// TAN: (101, 2) + /// TIT: (101, 0) + /// >>> print(table1.match_table(table2)) + /// [[101 5 100 4] + /// [101 0 100 3]] + #[pyo3(signature = (table, similarity_rule=None))] + fn match_table<'py>( + &self, + py: Python<'py>, + table: PyRef<'py, KmerTable>, + similarity_rule: Option>, + ) -> PyResult>> { + self.0.match_table(py, &table.0, similarity_rule.as_ref()) + } + + /// Find matches between the *k-mers* in this table with all + /// overlapping *k-mers* in the given `sequence`. + /// *k* is determined by the table. + /// + /// Parameters + /// ---------- + /// sequence : Sequence + /// The sequence to be matched. + /// The table's base alphabet must extend the alphabet of the + /// sequence. + /// similarity_rule : SimilarityRule, optional + /// If this parameter is given, not only exact *k-mer* matches + /// are considered, but also similar ones according to the given + /// :class:`SimilarityRule`. + /// ignore_mask : ndarray, dtype=bool, optional + /// Boolean mask of sequence positions to ignore. + /// *k-mers* that involve these sequence positions are not added + /// to the table. + /// This is used e.g. to skip repeat regions. + /// By default, no sequence position is ignored. + /// + /// Returns + /// ------- + /// matches : ndarray, shape=(n,3), dtype=np.int64 + /// The *k-mer* matches. + /// Each row contains one match. Each match has the following + /// columns: + /// + /// 0. The sequence position in the input sequence + /// 1. The reference ID of the matched sequence in the table + /// 2. The sequence position of the matched sequence in the + /// table + /// + /// Notes + /// ----- + /// + /// The matches are ordered by the first column. + /// + /// Examples + /// -------- + /// + /// >>> sequence1 = ProteinSequence("BIQTITE") + /// >>> table = KmerTable.from_sequences(3, [sequence1], ref_ids=[100]) + /// >>> print(table) + /// IQT: (100, 1) + /// ITE: (100, 4) + /// QTI: (100, 2) + /// TIT: (100, 3) + /// BIQ: (100, 0) + /// >>> sequence2 = ProteinSequence("TITANITE") + /// >>> print(table.match(sequence2)) + /// [[ 0 100 3] + /// [ 5 100 4]] + #[pyo3(name = "match", signature = (sequence, similarity_rule=None, ignore_mask=None))] + fn match_<'py>( + &self, + py: Python<'py>, + sequence: Bound<'py, PyAny>, + similarity_rule: Option>, + ignore_mask: Option>, + ) -> PyResult>> { + self.0.match_sequence( + py, + &sequence, + similarity_rule.as_ref(), + ignore_mask.as_ref(), + ) + } + + /// Find matches between the *k-mers* in this table with the given + /// *k-mer* selection. + /// + /// It is intended to use this method to find matches in a table + /// that was created using :meth:`from_kmer_selection()`. + /// + /// Parameters + /// ---------- + /// positions : ndarray, shape=(n,), dtype=uint32 + /// Sequence positions of the filtered subset of *k-mers* given + /// in `kmers`. + /// kmers : ndarray, shape=(n,), dtype=np.int64 + /// Filtered subset of *k-mer* codes to match against. + /// + /// Returns + /// ------- + /// matches : ndarray, shape=(n,3), dtype=np.int64 + /// The *k-mer* matches. + /// Each row contains one *k-mer* match. + /// Each match has the following columns: + /// + /// 0. The sequence position of the input *k-mer*, taken + /// from `positions` + /// 1. The reference ID of the matched sequence in the table + /// 2. The sequence position of the matched *k-mer* in the + /// table + fn match_kmer_selection<'py>( + &self, + py: Python<'py>, + positions: Bound<'py, PyAny>, + kmers: Bound<'py, PyAny>, + ) -> PyResult>> { + self.0.match_kmer_selection(py, &positions, &kmers) + } + + /// Count the number of occurences for each *k-mer* in the table. + /// + /// Parameters + /// ---------- + /// kmers : ndarray, dtype=np.int64, optional + /// The count is returned for these *k-mer* codes. + /// By default all *k-mers* are counted in ascending order, i.e. + /// ``count_for_kmer = counts[kmer]``. + /// + /// Returns + /// ------- + /// counts : ndarray, dtype=np.int64 + /// The counts for each given *k-mer*. + /// + /// Examples + /// -------- + /// >>> table = KmerTable.from_sequences( + /// ... k = 2, + /// ... sequences = [NucleotideSequence("TTATA"), NucleotideSequence("CTAG")], + /// ... ref_ids = [0, 1] + /// ... ) + /// >>> print(table.count(table.kmer_alphabet.encode_multiple(["TA", "AG"]))) + /// [3 1] + /// >>> counts = table.count() + /// >>> print(counts) + /// [0 0 1 1 0 0 0 1 0 0 0 0 3 0 0 1] + #[pyo3(signature = (kmers=None))] + fn count<'py>( + &self, + py: Python<'py>, + kmers: Option>, + ) -> PyResult>> { + match kmers { + Some(kmers) => self.0.count(py, &kmers), + None => Ok(self.0.count_all(py)), + } + } + + /// Get the *k-mer* codes for all *k-mers* that have at least one + /// position in the table. + /// + /// Returns + /// ------- + /// kmers : ndarray, shape=(n,), dtype=np.int64 + /// The *k-mer* codes in ascending order. + fn get_kmers<'py>(&self, py: Python<'py>) -> Bound<'py, PyArray1> { + self.0.get_kmers(py) + } + + fn __getitem__<'py>(&self, py: Python<'py>, kmer: Kmer) -> PyResult>> { + self.0.get_item(py, kmer) + } + + fn __len__(&self) -> usize { + self.0.n_kmers + } + + fn __contains__(&self, kmer: Kmer) -> bool { + self.0.contains(kmer) + } + + fn __iter__<'py>(&self, py: Python<'py>) -> PyResult> { + let kmers = PyList::new(py, self.0.collect_kmers())?; + Ok(kmers.try_iter()?.into_any()) + } + + fn __reversed__<'py>(&self, py: Python<'py>) -> PyResult> { + let mut kmers = self.0.collect_kmers(); + kmers.reverse(); + let kmers = PyList::new(py, kmers)?; + Ok(kmers.try_iter()?.into_any()) + } + + fn __eq__(&self, py: Python<'_>, other: Bound<'_, PyAny>) -> PyResult { + match other.extract::>() { + Ok(other) => self.0.equals(py, &other.0), + Err(_) => Ok(false), + } + } + + fn __str__(&self, py: Python<'_>) -> PyResult { + self.0.to_str(py) + } + + /// Pickle via a direct reconstructor, so unpickling does not build an empty + /// `n_slots`-sized table first. + fn __reduce__<'py>( + &self, + py: Python<'py>, + ) -> PyResult<(Bound<'py, PyAny>, Bound<'py, PyTuple>)> { + let reconstructor = py.get_type::().getattr("_from_state")?; + let (offsets, data) = self.0.pickle_state(py); + let args = PyTuple::new( + py, + [ + self.0.kmer_alphabet.clone_ref(py).into_bound(py), + offsets.into_any(), + data.into_any(), + ], + )?; + Ok((reconstructor, args)) + } +} + +/// This class represents a *k-mer* index table. +/// In contrast to :class:`KmerTable`, which gives every possible *k-mer* its +/// own slot, a :class:`BucketKmerTable` pools the *k-mers* into a limited +/// number of buckets. +/// Hence, different *k-mers* may be stored in the same bucket, like in a +/// hash table. +/// This approach makes *k-mer* indices with large *k-mer* alphabets +/// fit into memory. +/// +/// Otherwise, the API for creating a :class:`BucketKmerTable` and +/// matching to it is analogous to :class:`KmerTable`. +/// +/// Attributes +/// ---------- +/// kmer_alphabet : KmerAlphabet +/// The internal :class:`KmerAlphabet`, that is used to +/// encode all overlapping *k-mers* of an input sequence. +/// alphabet : Alphabet +/// The base alphabet, from which this :class:`BucketKmerTable` was +/// created. +/// k : int +/// The length of the *k-mers*. +/// n_buckets : int +/// The number of buckets, the *k-mers* are divided into. +/// +/// See Also +/// -------- +/// KmerTable +/// +/// Notes +/// ----- +/// +/// *Memory consumption* +/// +/// For efficient mapping, a :class:`BucketKmerTable` contains two large arrays: +/// +/// 1. An array that holds all *k-mers* with their positions. +/// 2. An accompanying array that maps each bucket to the index range in the +/// first array, where the corresponding entries are stored. +/// +/// As buckets are used, the memory requirements are limited to the number +/// of buckets instead of scaling with the :class:`KmerAlphabet` size. +/// If each bucket is used, the required memory space :math:`S` in byte +/// is +/// +/// .. math:: +/// +/// S = 8B + 16L +/// +/// where :math:`B` is the number of buckets and :math:`L` is the summed +/// length of all sequences added to the table. +/// +/// *Buckets* +/// +/// The ratio :math:`L/B` is called *load_factor*. +/// By default :class:`BucketKmerTable` uses a load factor of +/// approximately 0.8 to ensure efficient *k-mer* matching. +/// The number fo buckets can be adjusted by setting the +/// `n_buckets` parameters on :class:`BucketKmerTable` creation. +/// It is recommended to use :func:`bucket_number()` to compute an +/// appropriate number of buckets. +/// +/// *Multiprocessing* +/// +/// :class:`BucketKmerTable` objects can be used in multi-processed +/// setups: +/// Adding a large database of sequences to a table can be sped up by +/// splitting the database into smaller chunks and create a separate +/// table for each chunk in separate processes. +/// Eventually, the tables can be merged to one large table using +/// :meth:`from_tables()`. +/// +/// Since :class:`BucketKmerTable` supports the *pickle* protocol, +/// the matching step can also be divided into multiple processes, if +/// multiple sequences need to be matched. +/// +/// *Storage on hard drive* +/// +/// The most time efficient way to read/write a :class:`BucketKmerTable` +/// is the *pickle* format. +/// +/// *Indexing and iteration* +/// +/// Due to the higher complexity in the *k-mer* lookup compared to +/// :class:`KmerTable`, this class is still indexable but not iterable. +/// +/// Examples +/// -------- +/// +/// Create a *2-mer* index table for some nucleotide sequences: +/// +/// >>> table = BucketKmerTable.from_sequences( +/// ... k = 2, +/// ... sequences = [NucleotideSequence("TTATA"), NucleotideSequence("CTAG")], +/// ... ref_ids = [0, 1] +/// ... ) +/// +/// Display the contents of the table as +/// (reference ID, sequence position) tuples: +/// +/// >>> print(table) +/// AG: (1, 2) +/// AT: (0, 2) +/// CT: (1, 0) +/// TA: (0, 1), (0, 3), (1, 1) +/// TT: (0, 0) +/// +/// Find matches of the table with a sequence: +/// +/// >>> query = NucleotideSequence("TAG") +/// >>> matches = table.match(query) +/// >>> for query_pos, table_ref_id, table_pos in matches: +/// ... print("Query sequence position:", query_pos) +/// ... print("Table reference ID: ", table_ref_id) +/// ... print("Table sequence position:", table_pos) +/// ... print() +/// Query sequence position: 0 +/// Table reference ID: 0 +/// Table sequence position: 1 +/// +/// Query sequence position: 0 +/// Table reference ID: 0 +/// Table sequence position: 3 +/// +/// Query sequence position: 0 +/// Table reference ID: 1 +/// Table sequence position: 1 +/// +/// Query sequence position: 1 +/// Table reference ID: 1 +/// Table sequence position: 2 +/// +/// +/// Get all reference IDs and positions for a given *k-mer*: +/// +/// >>> kmer_code = table.kmer_alphabet.encode("TA") +/// >>> print(table[kmer_code]) +/// [[0 1] +/// [0 3] +/// [1 1]] +#[pyclass(module = "biotite.rust.sequence.align")] +pub struct BucketKmerTable(GenericKmerTable); + +#[pymethods] +impl BucketKmerTable { + /// Reconstruct a table from its pickled state (called by `__reduce__`). + #[staticmethod] + fn _from_state( + py: Python<'_>, + kmer_alphabet: KmerAlphabet, + offsets: PyReadonlyArray1<'_, i64>, + data: Bound<'_, PyBytes>, + ) -> PyResult { + Ok(BucketKmerTable(GenericKmerTable::from_state( + py, + kmer_alphabet, + offsets.as_slice()?, + data.as_bytes(), + )?)) + } + + #[getter] + fn kmer_alphabet(&self, py: Python<'_>) -> Py { + self.0.kmer_alphabet.clone_ref(py) + } + + #[getter] + fn alphabet<'py>(&self, py: Python<'py>) -> PyResult> { + self.0.kmer_alphabet.base_alphabet(py) + } + + #[getter] + fn k(&self) -> usize { + self.0.k + } + + #[getter] + fn n_buckets(&self) -> usize { + self.0.n_slots + } + + /// Create a :class:`BucketKmerTable` by storing the positions of + /// all overlapping *k-mers* from the input `sequences`. + /// + /// Parameters + /// ---------- + /// k : int + /// The length of the *k-mers*. + /// sequences : sized iterable object of Sequence, length=m + /// The sequences to get the *k-mer* positions from. + /// These sequences must have equal alphabets, or one of these + /// sequences must have an alphabet that extends the alphabets + /// of all other sequences. + /// ref_ids : sized iterable object of int, length=m, optional + /// The reference IDs for the given sequences. + /// By default the IDs are counted from *0* to *m*. + /// ignore_masks : sized iterable object of (ndarray, dtype=bool), length=m, optional + /// Sequence positions to ignore. + /// *k-mers* that involve these sequence positions are not added + /// to the table. + /// By default, no sequence position is ignored. + /// alphabet : Alphabet, optional + /// The alphabet to use for this table. + /// By default, an appropriate alphabet is inferred from the + /// input `sequences`. + /// spacing : None or str or list or ndarray, dtype=int, shape=(k,) + /// If provided, spaced *k-mers* are used instead of continuous + /// ones. + /// n_buckets : int, optional + /// Set the number of buckets in the table, e.g. to use a + /// different load factor. + /// It is recommended to use :func:`bucket_number()` for this + /// purpose. + /// By default, a load factor of approximately 0.8 is used. + /// + /// See Also + /// -------- + /// from_kmers : The same functionality based on already created *k-mers* + /// + /// Returns + /// ------- + /// table : BucketKmerTable + /// The newly created table. + /// + /// Examples + /// -------- + /// + /// >>> sequences = [NucleotideSequence("TTATA"), NucleotideSequence("CTAG")] + /// >>> table = BucketKmerTable.from_sequences( + /// ... 2, sequences, ref_ids=[100, 101] + /// ... ) + /// >>> print(table) + /// AG: (101, 2) + /// AT: (100, 2) + /// CT: (101, 0) + /// TA: (100, 1), (100, 3), (101, 1) + /// TT: (100, 0) + #[staticmethod] + #[pyo3(signature = (k, sequences, ref_ids=None, ignore_masks=None, alphabet=None, spacing=None, n_buckets=None))] + #[allow(clippy::too_many_arguments)] + fn from_sequences( + py: Python<'_>, + k: usize, + sequences: Vec>, + ref_ids: Option>, + ignore_masks: Option>>>, + alphabet: Option>, + spacing: Option>, + n_buckets: Option, + ) -> PyResult { + Ok(BucketKmerTable(GenericKmerTable::from_sequences( + py, + k, + sequences, + ref_ids, + ignore_masks, + alphabet, + spacing, + n_buckets, + )?)) + } + + /// Create a :class:`BucketKmerTable` by storing the positions of + /// all input *k-mers*. + /// + /// Parameters + /// ---------- + /// kmer_alphabet : KmerAlphabet + /// The :class:`KmerAlphabet` to use for the new table. + /// kmers : sized iterable object of (ndarray, dtype=np.int64), length=m + /// List where each array contains the *k-mer* codes from a + /// sequence. + /// ref_ids : sized iterable object of int, length=m, optional + /// The reference IDs for the sequences. + /// By default the IDs are counted from *0* to *m*. + /// masks : sized iterable object of (ndarray, dtype=bool), length=m, optional + /// A *k-mer* code at a position, where the corresponding mask + /// is false, is not added to the table. + /// By default, all positions are added. + /// n_buckets : int, optional + /// Set the number of buckets in the table. + /// By default, a load factor of approximately 0.8 is used. + /// + /// See Also + /// -------- + /// from_sequences : The same functionality based on undecomposed sequences + /// + /// Returns + /// ------- + /// table : BucketKmerTable + /// The newly created table. + /// + /// Examples + /// -------- + /// + /// >>> sequences = [ProteinSequence("BIQTITE"), ProteinSequence("NIQBITE")] + /// >>> kmer_alphabet = KmerAlphabet(ProteinSequence.alphabet, 3) + /// >>> kmer_codes = [kmer_alphabet.create_kmers(s.code) for s in sequences] + /// >>> table = BucketKmerTable.from_kmers(kmer_alphabet, kmer_codes) + /// >>> print(table) + /// IQT: (0, 1) + /// IQB: (1, 1) + /// ITE: (0, 4), (1, 4) + /// NIQ: (1, 0) + /// QTI: (0, 2) + /// QBI: (1, 2) + /// TIT: (0, 3) + /// BIQ: (0, 0) + /// BIT: (1, 3) + #[staticmethod] + #[pyo3(signature = (kmer_alphabet, kmers, ref_ids=None, masks=None, n_buckets=None))] + fn from_kmers( + py: Python<'_>, + kmer_alphabet: KmerAlphabet, + kmers: Vec>, + ref_ids: Option>, + masks: Option>>>, + n_buckets: Option, + ) -> PyResult { + Ok(BucketKmerTable(GenericKmerTable::from_kmers( + py, + kmer_alphabet, + kmers, + ref_ids, + masks, + n_buckets, + )?)) + } + + /// Create a :class:`BucketKmerTable` by storing the positions of a + /// filtered subset of input *k-mers*. + /// + /// This can be used to reduce the number of stored *k-mers* using + /// a *k-mer* subset selector such as :class:`MinimizerSelector`. + /// + /// Parameters + /// ---------- + /// kmer_alphabet : KmerAlphabet + /// The :class:`KmerAlphabet` to use for the new table. + /// positions : sized iterable object of (ndarray, shape=(n,), dtype=uint32), length=m + /// List where each array contains the sequence positions of + /// the filtered subset of *k-mers* given in `kmers`. + /// kmers : sized iterable object of (ndarray, shape=(n,), dtype=np.int64), length=m + /// List where each array contains the filtered subset of + /// *k-mer* codes from a sequence. + /// ref_ids : sized iterable object of int, length=m, optional + /// The reference IDs for the sequences. + /// By default the IDs are counted from *0* to *m*. + /// n_buckets : int, optional + /// Set the number of buckets in the table. + /// By default, a load factor of approximately 0.8 is used. + /// + /// Returns + /// ------- + /// table : BucketKmerTable + /// The newly created table. + #[staticmethod] + #[pyo3(signature = (kmer_alphabet, positions, kmers, ref_ids=None, n_buckets=None))] + fn from_kmer_selection( + py: Python<'_>, + kmer_alphabet: KmerAlphabet, + positions: Vec>, + kmers: Vec>, + ref_ids: Option>, + n_buckets: Option, + ) -> PyResult { + Ok(BucketKmerTable(GenericKmerTable::from_kmer_selection( + py, + kmer_alphabet, + positions, + kmers, + ref_ids, + n_buckets, + )?)) + } + + /// Create a :class:`BucketKmerTable` by merging the *k-mer* + /// positions from existing `tables`. + /// + /// Parameters + /// ---------- + /// tables : iterable object of BucketKmerTable + /// The tables to be merged. + /// All tables must have equal number of buckets and equal + /// :class:`KmerAlphabet` objects, i.e. the same *k* and equal + /// base alphabets. + /// + /// Returns + /// ------- + /// table : BucketKmerTable + /// The newly created table. + #[staticmethod] + fn from_tables(py: Python<'_>, tables: Vec>) -> PyResult { + let refs: Vec<&GenericKmerTable> = tables.iter().map(|table| &table.0).collect(); + Ok(BucketKmerTable(GenericKmerTable::from_tables(py, &refs)?)) + } + + /// Find matches between the *k-mers* in this table with the + /// *k-mers* in another `table`. + /// + /// Parameters + /// ---------- + /// table : BucketKmerTable + /// The table to be matched. + /// Both tables must have equal number of buckets and equal + /// :class:`KmerAlphabet` objects. + /// similarity_rule : SimilarityRule, optional + /// If this parameter is given, not only exact *k-mer* matches + /// are considered, but also similar ones according to the given + /// :class:`SimilarityRule`. + /// + /// Returns + /// ------- + /// matches : ndarray, shape=(n,4), dtype=np.int64 + /// The *k-mer* matches. + /// + /// Notes + /// ----- + /// + /// There is no guaranteed order of the reference IDs or + /// sequence positions in the returned matches. + #[pyo3(signature = (table, similarity_rule=None))] + fn match_table<'py>( + &self, + py: Python<'py>, + table: PyRef<'py, BucketKmerTable>, + similarity_rule: Option>, + ) -> PyResult>> { + self.0.match_table(py, &table.0, similarity_rule.as_ref()) + } + + /// Find matches between the *k-mers* in this table with all + /// overlapping *k-mers* in the given `sequence`. + /// + /// Parameters + /// ---------- + /// sequence : Sequence + /// The sequence to be matched. + /// similarity_rule : SimilarityRule, optional + /// If this parameter is given, not only exact *k-mer* matches + /// are considered, but also similar ones according to the given + /// :class:`SimilarityRule`. + /// ignore_mask : ndarray, dtype=bool, optional + /// Boolean mask of sequence positions to ignore. + /// By default, no sequence position is ignored. + /// + /// Returns + /// ------- + /// matches : ndarray, shape=(n,3), dtype=np.int64 + /// The *k-mer* matches. + /// + /// Notes + /// ----- + /// + /// The matches are ordered by the first column. + #[pyo3(name = "match", signature = (sequence, similarity_rule=None, ignore_mask=None))] + fn match_<'py>( + &self, + py: Python<'py>, + sequence: Bound<'py, PyAny>, + similarity_rule: Option>, + ignore_mask: Option>, + ) -> PyResult>> { + self.0.match_sequence( + py, + &sequence, + similarity_rule.as_ref(), + ignore_mask.as_ref(), + ) + } + + /// Find matches between the *k-mers* in this table with the given + /// *k-mer* selection. + /// + /// Parameters + /// ---------- + /// positions : ndarray, shape=(n,), dtype=uint32 + /// Sequence positions of the filtered subset of *k-mers* given + /// in `kmers`. + /// kmers : ndarray, shape=(n,), dtype=np.int64 + /// Filtered subset of *k-mer* codes to match against. + /// + /// Returns + /// ------- + /// matches : ndarray, shape=(n,3), dtype=np.int64 + /// The *k-mer* matches. + fn match_kmer_selection<'py>( + &self, + py: Python<'py>, + positions: Bound<'py, PyAny>, + kmers: Bound<'py, PyAny>, + ) -> PyResult>> { + self.0.match_kmer_selection(py, &positions, &kmers) + } + + /// Count the number of occurences for each given *k-mer* in the + /// table. + /// + /// Parameters + /// ---------- + /// kmers : ndarray, dtype=np.int64 + /// The count is returned for these *k-mer* codes. + /// + /// Returns + /// ------- + /// counts : ndarray, dtype=np.int64 + /// The counts for each given *k-mer*. + /// + /// Notes + /// ----- + /// As each bucket need to be inspected for the actual *k-mer* + /// entries, this method requires far more computation time than its + /// :class:`KmerTable` equivalent. + fn count<'py>( + &self, + py: Python<'py>, + kmers: Bound<'py, PyAny>, + ) -> PyResult>> { + self.0.count(py, &kmers) + } + + /// Get the *k-mer* codes for all *k-mers* that have at least one + /// position in the table. + /// + /// Returns + /// ------- + /// kmers : ndarray, shape=(n,), dtype=np.int64 + /// The *k-mer* codes in ascending order. + /// + /// Notes + /// ----- + /// As each bucket need to be inspected for the actual *k-mer* + /// entries, this method requires far more computation time than its + /// :class:`KmerTable` equivalent. + fn get_kmers<'py>(&self, py: Python<'py>) -> Bound<'py, PyArray1> { + self.0.get_kmers(py) + } + + fn __getitem__<'py>(&self, py: Python<'py>, kmer: Kmer) -> PyResult>> { + self.0.get_item(py, kmer) + } + + fn __len__(&self) -> usize { + self.0.n_kmers + } + + fn __eq__(&self, py: Python<'_>, other: Bound<'_, PyAny>) -> PyResult { + match other.extract::>() { + Ok(other) => self.0.equals(py, &other.0), + Err(_) => Ok(false), + } + } + + fn __str__(&self, py: Python<'_>) -> PyResult { + self.0.to_str(py) + } + + /// Pickle via a direct reconstructor, so unpickling does not build an empty + /// `n_slots`-sized table first. + fn __reduce__<'py>( + &self, + py: Python<'py>, + ) -> PyResult<(Bound<'py, PyAny>, Bound<'py, PyTuple>)> { + let reconstructor = py.get_type::().getattr("_from_state")?; + let (offsets, data) = self.0.pickle_state(py); + let args = PyTuple::new( + py, + [ + self.0.kmer_alphabet.clone_ref(py).into_bound(py), + offsets.into_any(), + data.into_any(), + ], + )?; + Ok((reconstructor, args)) + } +} diff --git a/src/rust/sequence/align/mod.rs b/src/rust/sequence/align/mod.rs index ca495d20a..18f102af7 100644 --- a/src/rust/sequence/align/mod.rs +++ b/src/rust/sequence/align/mod.rs @@ -2,8 +2,10 @@ use pyo3::prelude::*; pub mod banded; pub mod cell; +pub mod kmertable; pub mod localgapped; pub mod localungapped; +pub mod nested; pub mod pairwise; pub mod scoring; pub mod symbol; @@ -16,5 +18,7 @@ pub fn module<'py>(parent_module: &Bound<'py, PyModule>) -> PyResult()?; + module.add_class::()?; + module.add_class::()?; Ok(module) } diff --git a/src/rust/sequence/align/nested.rs b/src/rust/sequence/align/nested.rs new file mode 100644 index 000000000..3fc03064a --- /dev/null +++ b/src/rust/sequence/align/nested.rs @@ -0,0 +1,159 @@ +use std::ops::{Index, IndexMut}; + +/// A frozen container representing a nested array of fixed-length inner arrays. +/// +/// All inner arrays are stored back-to-back in a single flat `data` buffer, with +/// an `offsets` array recording where each one begins and ends. +/// Use index operations to access inner arrays. +pub struct NestedArray { + /// The flat buffer holding every inner array consecutively; inner array `i` + /// occupies ``data[offsets[i]..offsets[i + 1]]``. + data: Vec, + /// The bounds of each inner array within `data`. Has length + /// ``n_inner_arrays + 1``, is non-decreasing, starts at ``0`` and ends at + /// ``data.len()``, so inner array `i` spans ``offsets[i]..offsets[i + 1]``. + offsets: Vec, +} + +impl NestedArray { + /// Create a `NestedArray` directly from the flat `data` buffer and the + /// `offsets` array, i.e. the inverse of `NestedArray::raw_parts`. + /// + /// Safety + /// ------ + /// `offsets` must have length ``n_inner_arrays + 1``, be non-decreasing, + /// start at ``0`` and end at ``data.len()``. + pub unsafe fn from_raw_parts(data: Vec, offsets: Vec) -> Self { + debug_assert!(!offsets.is_empty()); + debug_assert_eq!(offsets[0], 0); + debug_assert_eq!(*offsets.last().unwrap(), data.len()); + Self { data, offsets } + } + + /// Borrow the flat data buffer and the offsets array, i.e. the inverse of + /// `NestedArray::from_raw_parts`. + /// + /// `offsets` has length ``n_inner_arrays + 1``; inner array `i` occupies + /// ``data[offsets[i]..offsets[i + 1]]``. + /// Can be used to serialize the table. + pub fn raw_parts(&self) -> (&[T], &[usize]) { + (&self.data, &self.offsets) + } +} + +impl Index for NestedArray { + type Output = [T]; + fn index(&self, index: usize) -> &Self::Output { + &self.data[self.offsets[index]..self.offsets[index + 1]] + } +} + +impl IndexMut for NestedArray { + fn index_mut(&mut self, index: usize) -> &mut Self::Output { + &mut self.data[self.offsets[index]..self.offsets[index + 1]] + } +} + +/// Builder for a `NestedArray`: The inner array lengths are fixed up +/// front and elements are scattered into their slots via `push`. +/// +/// Notes +/// ----- +/// The per-slot write cursors are stored in the `offsets` array itself (no +/// separate cursor array), and the data buffer is filled while uninitialized +/// (no zero-initialization). Both are key to construction performance for large +/// tables. +/// +/// ``T: Copy`` is required because the elements are written through a raw +/// pointer while the backing `Vec`'s length stays ``0`` until `build`. If the +/// builder were dropped before `build` (e.g. on a panic mid-fill), those +/// elements would not be seen by the `Vec` and thus not dropped. Forbidding +/// `Drop` (which `Copy` does) means there is nothing to leak. +pub struct NestedArrayBuilder { + /// The flat data buffer. `Vec::len` stays ``0`` (the buffer is + /// uninitialized) until `build` sets it to `expected_total`; + /// elements are written through the raw pointer. + data: Vec, + /// Per-slot write cursor: ``offsets[i]`` is the next free position of slot + /// `i`. `build` un-shifts it into the final slot-start offsets. + offsets: Vec, + /// The total number of elements across all slots (the data capacity). + expected_total: usize, + /// The number of elements added so far via `push`. + elements_added: usize, +} + +impl NestedArrayBuilder { + /// Create a builder for a `NestedArray` whose inner array ``i`` holds exactly + /// ``lengths[i]`` elements. + pub fn new(lengths: Vec) -> Self { + // Prefix-sum the lengths into slot start offsets, reusing the vector + let mut offsets = lengths; + let mut total = 0; + for slot in &mut offsets { + let length = *slot; + *slot = total; + total += length; + } + offsets.push(total); + Self { + data: Vec::with_capacity(total), + offsets, + expected_total: total, + elements_added: 0, + } + } + + /// Append `value` to inner array `index`. + /// + /// Safety + /// ------ + /// Across all calls, the number of elements pushed to a given `index` must + /// not exceed that slot's length (as passed to `new`). Otherwise a write + /// would spill into a neighboring slot's region and leave another slot + /// partly uninitialized, which `build` would then expose. + #[inline] + pub unsafe fn push(&mut self, index: usize, value: T) { + assert!( + self.elements_added < self.expected_total, + "pushed more elements than the total slot length" + ); + let cursor = &mut self.offsets[index]; + let position = *cursor; + *cursor = position + 1; + // SAFETY: by the per-slot contract, `position` stays within slot + // `index`, hence `position < expected_total == capacity`, and each + // position is written exactly once. + unsafe { + self.data.as_mut_ptr().add(position).write(value); + } + self.elements_added += 1; + } + + /// Finalize the builder into a `NestedArray`. + /// + /// Panics + /// ------ + /// Panics if fewer elements were added than the total slot length, since + /// that would leave part of the data buffer uninitialized. + pub fn build(mut self) -> NestedArray { + assert_eq!( + self.elements_added, self.expected_total, + "fewer elements were added than the total slot length" + ); + // SAFETY: exactly `expected_total` elements were written, one per + // position in `0..expected_total`. + unsafe { + self.data.set_len(self.expected_total); + } + // Each `offsets[i]` has advanced to the start of slot `i + 1`; + // shift right by one and restore `offsets[0] = 0` to recover the slot starts + let n_slots = self.offsets.len() - 1; + self.offsets.copy_within(0..n_slots, 1); + self.offsets[0] = 0; + // SAFETY: `offsets` is the prefix-sum offset array (length `n_slots + 1`, + // non-decreasing, starting at `0` and ending at `expected_total`), and + // `data` was just filled with exactly `expected_total` elements. + unsafe { NestedArray::from_raw_parts(self.data, self.offsets) } + } +} diff --git a/src/rust/structure/bonds.rs b/src/rust/structure/bonds.rs index 9bda8a81d..2c82b702a 100644 --- a/src/rust/structure/bonds.rs +++ b/src/rust/structure/bonds.rs @@ -148,8 +148,6 @@ impl Bond { } } -/// __init__(atom_count, bonds=None) -/// /// A bond list stores indices of atoms /// (usually of an :class:`AtomArray` or :class:`AtomArrayStack`) /// that form chemical bonds together with the type (or order) of the @@ -420,8 +418,6 @@ impl BondList { BondList::concatenate_lists(&refs) } - /// offset_indices(offset) - /// /// Increase all atom indices in the :class:`BondList` by the given /// offset. /// @@ -459,8 +455,6 @@ impl BondList { Ok(()) } - /// as_array() - /// /// Obtain a copy of the internal :class:`ndarray`. /// /// Returns @@ -520,8 +514,6 @@ impl BondList { Ok(set) } - /// as_graph() - /// /// Obtain a graph representation of the :class:`BondList`. /// /// Returns @@ -623,8 +615,6 @@ impl BondList { } } - /// convert_bond_type(original_bond_type, new_bond_type) - /// /// Convert all occurences of a given bond type into another bond type. /// /// Parameters @@ -681,8 +671,6 @@ impl BondList { self.bonds.len() } - /// get_bonds(atom_index) - /// /// Obtain the indices of the atoms bonded to the atom with the /// given index as well as the corresponding bond types. /// @@ -963,8 +951,6 @@ impl BondList { Ok(()) } - /// remove_bond(atom_index1, atom_index2) - /// /// Remove a bond from the :class:`BondList`. /// /// If the bond is not existent in the :class:`BondList`, nothing happens. @@ -979,8 +965,6 @@ impl BondList { Ok(()) } - /// remove_bonds_to(self, atom_index) - /// /// Remove all bonds from the :class:`BondList` where the given atom /// is involved. /// @@ -995,8 +979,6 @@ impl BondList { Ok(()) } - /// remove_bonds(bond_list) - /// /// Remove multiple bonds from the :class:`BondList`. /// /// All bonds present in `bond_list` are removed from this instance. @@ -1015,8 +997,6 @@ impl BondList { .retain(|bond| !bonds_to_remove.contains(&(bond.atom1, bond.atom2))); } - /// merge(bond_list) - /// /// Merge another :class:`BondList` with this instance into a new /// object. /// If a bond appears in both :class:`BondList`'s, the diff --git a/src/rust/structure/celllist.rs b/src/rust/structure/celllist.rs index bfdd855b5..49e2c8c2a 100644 --- a/src/rust/structure/celllist.rs +++ b/src/rust/structure/celllist.rs @@ -296,8 +296,6 @@ impl IndexMut<[isize; 3]> for CellGrid { } } -/// __init__(atom_array, cell_size, periodic=False, box=None, selection=None) -/// /// This class enables the efficient search of atoms in vicinity of a /// defined location. /// diff --git a/src/rust/structure/connect.rs b/src/rust/structure/connect.rs index 5209112b4..3969b7a24 100644 --- a/src/rust/structure/connect.rs +++ b/src/rust/structure/connect.rs @@ -254,8 +254,6 @@ pub fn connect_inter_residue<'py>( Ok(bond_list) } -/// find_connected(bond_list, root, as_mask=False) -/// /// Get indices to all atoms that are directly or indirectly connected /// to the root atom indicated by the given index. /// diff --git a/src/rust/structure/io/pdb/file.rs b/src/rust/structure/io/pdb/file.rs index 470aa7a9e..3df6108dd 100644 --- a/src/rust/structure/io/pdb/file.rs +++ b/src/rust/structure/io/pdb/file.rs @@ -92,8 +92,6 @@ impl PDBFile { Ok(PyArray1::from_iter(py, self.atom_line_i.iter().map(|x| *x as i64)).to_owned()) } - /// read(file) - /// /// Read a PDB file. /// /// Parameters @@ -153,8 +151,6 @@ impl PDBFile { Ok(()) } - /// get_remark(self, number) - /// /// Get the lines containing the *REMARK* records with the given /// `number`. /// @@ -215,8 +211,6 @@ impl PDBFile { } } - /// get_model_count(self) - /// /// Get the number of models contained in the PDB file. /// /// Returns @@ -227,8 +221,6 @@ impl PDBFile { self.model_start_i.len() } - /// get_coord(self, model=None) - /// /// Get only the coordinates from the PDB file. /// /// Parameters @@ -339,8 +331,6 @@ impl PDBFile { } } - /// get_b_factor(self, model=None) - /// /// Get only the B-factors from the PDB file. /// /// Parameters @@ -701,8 +691,6 @@ impl PDBFile { Ok(atoms) } - /// set_structure(self, atoms, hybrid36=False) - /// /// Set the :class:`AtomArray` or :class:`AtomArrayStack` for the /// file. /// diff --git a/src/rust/structure/io/pdb/hybrid36.rs b/src/rust/structure/io/pdb/hybrid36.rs index 5e2759536..bbe620dc9 100644 --- a/src/rust/structure/io/pdb/hybrid36.rs +++ b/src/rust/structure/io/pdb/hybrid36.rs @@ -118,7 +118,6 @@ pub fn decode_hybrid36(hybrid36: &str) -> PyResult { } } -/// max_hybrid36_number(length) /// -- /// /// Give the maximum integer value that can be represented by a diff --git a/tests/sequence/align/test_kmertable.py b/tests/sequence/align/test_kmertable.py index 1ed10d37f..9c7cc4e30 100644 --- a/tests/sequence/align/test_kmertable.py +++ b/tests/sequence/align/test_kmertable.py @@ -50,16 +50,16 @@ def alphabet(): return seq.NucleotideSequence.unambiguous_alphabet() -@pytest.fixture -def random_sequences(k, alphabet): +@pytest.fixture(params=range(3)) +def random_sequences(request, k, alphabet): N_SEQS = 10 SEQ_LENGTH = 1000 - np.random.seed(0) + rng = np.random.default_rng(request.param) sequences = [] for _ in range(N_SEQS): sequence = seq.NucleotideSequence() - sequence.code = np.random.randint(len(alphabet), size=SEQ_LENGTH) + sequence.code = rng.integers(len(alphabet), size=SEQ_LENGTH) sequences.append(sequence) return sequences @@ -82,7 +82,7 @@ def random_sequences(k, alphabet): def test_from_sequences(k, random_sequences, spacing, table_class): """ Test the :meth:`from_sequences()` constructor, by checking for each - sequence position, if the position is in the C-array of the + sequence position, if the position is in the internal array of the corresponding k-mer. """ table = table_class.from_sequences(k, random_sequences, spacing=spacing) @@ -136,9 +136,9 @@ def test_from_kmer_selection(k, alphabet, random_sequences, table_class): kmer_arrays = [ kmer_alph.create_kmers(sequence.code) for sequence in random_sequences ] - np.random.seed(0) + rng = np.random.default_rng(0) filtered_pos_arrays = [ - np.random.randint(len(kmers), size=N_POSITIONS) for kmers in kmer_arrays + rng.integers(len(kmers), size=N_POSITIONS) for kmers in kmer_arrays ] filtered_kmer_arrays = [ kmers[filtered_pos] @@ -305,8 +305,8 @@ def test_match_kmer_selection(k, random_sequences, table_class): table = table_class.from_sequences(k, table_sequences) kmers = table.kmer_alphabet.create_kmers(query_sequence.code) - np.random.seed(0) - positions = np.random.randint(len(kmers), size=N_POS) + rng = np.random.default_rng(0) + positions = rng.integers(len(kmers), size=N_POS) ref_matches = [] for pos in positions: kmer = kmers[pos] @@ -341,9 +341,9 @@ def test_match_equivalence(k, random_sequences, table_class, use_mask): if use_mask: # Create a random removal mask with low removal probability - np.random.seed(0) + rng = np.random.default_rng(0) removal_masks = [ - np.random.choice([False, True], size=len(sequence), p=[0.95, 0.05]) + rng.choice([False, True], size=len(sequence), p=[0.95, 0.05]) for sequence in random_sequences ] else: @@ -432,8 +432,8 @@ def test_count(k, random_sequences, table_class, selected_kmers): table = table_class.from_sequences(k, random_sequences) if selected_kmers: - np.random.seed(0) - kmers = np.random.randint(len(table.kmer_alphabet), size=N_KMERS) + rng = np.random.default_rng(0) + kmers = rng.integers(len(table.kmer_alphabet), size=N_KMERS) ref_counts = [len(table[kmer]) for kmer in kmers] test_counts = table.count(kmers) else: @@ -454,10 +454,10 @@ def test_get_kmers(table_class): :meth:`get_kmers()`, by constructing a table with exactly one appearance for each *k-mer* in a random list of *k-mers*. """ - np.random.seed(0) + rng = np.random.default_rng(0) kmer_alphabet = align.KmerAlphabet(seq.NucleotideSequence.unambiguous_alphabet(), 8) - ref_mask = np.random.choice([False, True], size=len(kmer_alphabet)) + ref_mask = rng.choice([False, True], size=len(kmer_alphabet)) ref_kmers = np.where(ref_mask)[0] table = table_class.from_kmers(kmer_alphabet, [ref_kmers])