kaptive.core.kmers¶

def __len__(self) -> int:
    r"""Return the number of records in the index.

    Returns:
        int: Record count.
    """
    return len(self.records)

@classmethod
def build(cls, batch: Sequences, **kwargs: Any) -> BaseKmerIndex:
    r"""Build a k-mer index from a sequence collection.

    Args:
        batch (Sequences): Sequence collection to index.
        **kwargs (Any): Subclass-specific options.

    Raises:
        NotImplementedError: Must be implemented by subclasses.
    """
    raise NotImplementedError

@classmethod
def empty(cls) -> BaseKmerIndex:
    r"""Create an empty BaseKmerIndex.

    Raises:
        NotImplementedError: Must be implemented by subclasses.
    """
    raise NotImplementedError

def top_hits(self, queries: BaseKmerIndex | Sequences, min_score: int = 1) -> Seeds:
    r"""Find the single best-matching target sequence for each query sequence.

    Args:
        queries (BaseKmerIndex | Sequences): Query sequence index or raw sequences.
        min_score (int): Minimum match score threshold. Defaults to 1.

    Returns:
        Seeds: Matched [`Seeds`][kaptive.core.kmers.Seeds] collection.
    """
    if len(queries) == 0 or len(self) == 0:
        return Seeds.empty()

    queries_idx = self._prep_queries(queries)
    q_offsets = _compute_query_offsets(queries_idx.records, queries_idx.n_seqs)

    seeds = Seeds(
        *_intersect_top_hit_kernel(queries_idx.records, q_offsets, queries_idx.n_seqs, self.records, self.n_seqs)
    )

    if min_score > 0:
        return seeds.filter(seeds.scores >= min_score)

    return seeds

def __len__(self) -> int:
    r"""Return the number of records in the index.

    Returns:
        int: Record count.
    """
    return len(self.records)

@classmethod
def build(
    cls,
    batch: Sequences,
    k: int = 21,
    scaled: int = 100,
    canonical: bool = True,
    seed: int = 42,
    sort_by_hash: bool = False,
    lut: npt.NDArray[np.uint8] | None = None,
    bits_per_char: int = 2,
    **kwargs: Any,
) -> FracMinHashIndex:
    r"""Build a FracMinHashIndex from nucleotide sequence batch.

    Args:
        batch (Sequences): Input sequence collection.
        k (int): K-mer length. Defaults to 21.
        scaled (int): Sampling scale factor. Defaults to 100.
        canonical (bool): True if using canonical k-mers. Defaults to True.
        seed (int): Seed for hashing. Defaults to 42.
        sort_by_hash (bool): True to sort output records by hash. Defaults to False.
        lut (npt.NDArray[np.uint8] | None): Optional DNA lookup table.
        bits_per_char (int): Bits per character (2 for DNA).
        **kwargs (Any): Additional parameters.

    Returns:
        FracMinHashIndex: Built [`FracMinHashIndex`][kaptive.core.kmers.FracMinHashIndex] instance.
    """
    if len(batch) == 0:
        return cls.empty()

    n_seqs = len(batch)
    kernel_lut = lut if lut is not None else _dna_lut()

    # Pass 1: Count fracminhashes per sequence to allocate exact memory footprint
    counts = _count_fracminhash_kernel(
        batch.seqs, batch.offsets, batch.lengths, kernel_lut, k, scaled, canonical, bits_per_char
    )

    if (total_hashes := np.sum(counts)) == 0:
        return cls.empty()

    # Generate exact write offsets via cumulative sum
    out_offsets = np.empty(n_seqs, dtype=np.uint32)
    current_offset = 0
    for i in range(n_seqs):
        out_offsets[i] = current_offset
        current_offset += counts[i]

    out_records = np.empty(total_hashes, dtype=MINHASH_DTYPE)

    # Pass 2: Populate the fracminhashes in parallel
    _populate_fracminhash_kernel(
        batch.seqs,
        batch.offsets,
        batch.lengths,
        kernel_lut,
        k,
        scaled,
        canonical,
        bits_per_char,
        out_offsets,
        out_records,
    )

    if sort_by_hash:
        out_records = _radix_sort_records(out_records)

    return cls(
        records=out_records,
        n_seqs=n_seqs,
        is_sorted=sort_by_hash,
        k=k,
        scaled=scaled,
        canonical=canonical,
        bits_per_char=bits_per_char,
        lut=lut,
    )

@classmethod
def empty(cls) -> FracMinHashIndex:
    r"""Create an empty FracMinHashIndex.

    Returns:
        FracMinHashIndex: Empty index instance.
    """
    return cls(
        records=np.empty(0, dtype=MINHASH_DTYPE),
        n_seqs=0,
        is_sorted=False,
        k=21,
        scaled=100,
        canonical=True,
        bits_per_char=2,
        lut=None,
    )

def to_sorted(self) -> FracMinHashIndex:
    r"""Return a new FracMinHashIndex with records sorted by hash.

    Returns:
        FracMinHashIndex: Sorted [`FracMinHashIndex`][kaptive.core.kmers.FracMinHashIndex].
    """
    if self.is_sorted:
        return self
    return self.__class__(
        records=_radix_sort_records(self.records),
        n_seqs=self.n_seqs,
        is_sorted=True,
        k=self.k,
        scaled=self.scaled,
        canonical=self.canonical,
        bits_per_char=self.bits_per_char,
        lut=self.lut,
    )

def top_hits(self, queries: BaseKmerIndex | Sequences, min_score: int = 1) -> Seeds:
    r"""Find the single best-matching target sequence for each query sequence.

    Args:
        queries (BaseKmerIndex | Sequences): Query sequence index or raw sequences.
        min_score (int): Minimum match score threshold. Defaults to 1.

    Returns:
        Seeds: Matched [`Seeds`][kaptive.core.kmers.Seeds] collection.
    """
    if len(queries) == 0 or len(self) == 0:
        return Seeds.empty()

    queries_idx = self._prep_queries(queries)
    q_offsets = _compute_query_offsets(queries_idx.records, queries_idx.n_seqs)

    seeds = Seeds(
        *_intersect_top_hit_kernel(queries_idx.records, q_offsets, queries_idx.n_seqs, self.records, self.n_seqs)
    )

    if min_score > 0:
        return seeds.filter(seeds.scores >= min_score)

    return seeds

def __len__(self) -> int:
    r"""Return the number of records in the index.

    Returns:
        int: Record count.
    """
    return len(self.records)

@classmethod
def build(
    cls,
    batch: Sequences,
    k: int = 10,
    s: int = 5,
    w_min: int = 1,
    w_max: int = 5,
    canonical: bool = True,
    seed: int = 42,
    sort_by_hash: bool = False,
    lut: npt.NDArray[np.uint8] | None = None,
    **kwargs: Any,
) -> RandstrobeIndex:
    r"""Build a RandstrobeIndex from sequence batch.

    Args:
        batch (Sequences): Input sequence collection.
        k (int): K-mer length. Defaults to 10.
        s (int): Sub-k-mer size. Defaults to 5.
        w_min (int): Min syncmer window bound. Defaults to 1.
        w_max (int): Max syncmer window bound. Defaults to 5.
        canonical (bool): Canonical option. Defaults to True.
        seed (int): Random seed. Defaults to 42.
        sort_by_hash (bool): True to sort records by hash. Defaults to False.
        lut (npt.NDArray[np.uint8] | None): Optional alphabet lookup table.
        **kwargs (Any): Additional parameters.

    Returns:
        RandstrobeIndex: Constructed [`RandstrobeIndex`][kaptive.core.kmers.RandstrobeIndex].

    Raises:
        ValueError: If `s >= k`.
    """
    if s >= k:
        raise ValueError("Sub-k-mer size (s) must be strictly less than k-mer size (k).")

    if len(batch) == 0:
        return cls.empty()

    # Pass 1: Count randstrobes per sequence to allocate exact memory footprint
    kernel_lut = lut if lut is not None else _mmseqs12_lut()
    counts = _count_randstrobes_kernel(batch.seqs, batch.offsets, batch.lengths, kernel_lut, k, s, w_min)
    if (total_randstrobes := np.sum(counts)) == 0:
        return cls.empty()

    # Generate exact write offsets via cumulative sum
    out_offsets = np.empty(len(counts), dtype=np.uint32)
    current_offset = 0
    for i in range(len(counts)):
        out_offsets[i] = current_offset
        current_offset += counts[i]

    # Pre-allocate output arrays
    out_records = np.empty(total_randstrobes, dtype=RANDSTROBE_DTYPE)

    # Pass 2: Populate the randstrobes in parallel
    _populate_randstrobes_kernel(
        batch.seqs, batch.offsets, batch.lengths, kernel_lut, k, s, w_min, w_max, out_offsets, out_records
    )

    if sort_by_hash:
        out_records = _radix_sort_records(out_records)

    return cls(
        records=out_records,
        n_seqs=len(batch),
        is_sorted=sort_by_hash,
        k=k,
        s=s,
        w_min=w_min,
        w_max=w_max,
        lut=lut,
    )

@classmethod
def empty(cls) -> RandstrobeIndex:
    r"""Create an empty RandstrobeIndex.

    Returns:
        RandstrobeIndex: Empty index instance.
    """
    return cls(
        records=np.empty(0, dtype=RANDSTROBE_DTYPE),
        n_seqs=0,
        is_sorted=False,
        k=10,
        s=5,
        w_min=1,
        w_max=5,
        lut=None,
    )

def top_hits(self, queries: BaseKmerIndex | Sequences, min_score: int = 1) -> Seeds:
    r"""Find the single best-matching target sequence for each query sequence.

    Args:
        queries (BaseKmerIndex | Sequences): Query sequence index or raw sequences.
        min_score (int): Minimum match score threshold. Defaults to 1.

    Returns:
        Seeds: Matched [`Seeds`][kaptive.core.kmers.Seeds] collection.
    """
    if len(queries) == 0 or len(self) == 0:
        return Seeds.empty()

    queries_idx = self._prep_queries(queries)
    q_offsets = _compute_query_offsets(queries_idx.records, queries_idx.n_seqs)

    seeds = Seeds(
        *_intersect_top_hit_kernel(queries_idx.records, q_offsets, queries_idx.n_seqs, self.records, self.n_seqs)
    )

    if min_score > 0:
        return seeds.filter(seeds.scores >= min_score)

    return seeds

def __getitem__(self, item: Any) -> Seed | Seeds:
    r"""Access seeds by integer index, slice, or boolean/integer NumPy array.

    Args:
        item (Any): Integer index, slice, or mask array.

    Returns:
        Seed | Seeds: A single [`Seed`][kaptive.core.kmers.Seed] or a
            sliced [`Seeds`][kaptive.core.kmers.Seeds] collection.

    Raises:
        IndexError: If integer index is out of bounds.
    """
    if isinstance(item, (int, np.integer)):
        if item < 0:
            item += len(self)
        if item < 0 or item >= len(self):
            raise IndexError("Batch index out of range")
        return Seed(
            int(self.query_indices[item]),
            int(self.target_indices[item]),
            int(self.scores[item]),
            int(self.offsets[item]),
        )

    if isinstance(item, slice):
        indices = np.arange(len(self))[item]
    else:
        item_arr = np.asarray(item)
        indices = np.nonzero(item_arr)[0] if item_arr.dtype.kind == "b" else item_arr

    return Seeds(
        self.query_indices[indices],
        self.target_indices[indices],
        self.scores[indices],
        self.offsets[indices],
    )

def __len__(self) -> int:
    r"""Return the number of seeds in the batch.

    Returns:
        int: Total seed count.
    """
    return len(self.query_indices)

@classmethod
def concat(cls, batches: Iterable[Self]) -> Self:  # type: ignore
    r"""Concatenate multiple [`Seeds`][kaptive.core.kmers.Seeds] collections into a single batch.

    Args:
        batches (Iterable[Seeds]): Iterable of [`Seeds`][kaptive.core.kmers.Seeds] objects.

    Returns:
        Seeds: Concatenated [`Seeds`][kaptive.core.kmers.Seeds] collection.
    """
    batches_list = list(batches)
    if not batches_list:
        return cls.empty()  # type: ignore
    return cls(
        np.concatenate([b.query_indices for b in batches_list]),
        np.concatenate([b.target_indices for b in batches_list]),
        np.concatenate([b.scores for b in batches_list]),
        np.concatenate([b.offsets for b in batches_list]),
    )

def cull_overlaps(
    self,
    query_lengths: npt.NDArray[np.int32],
    max_overlap_fraction: float = 0.1,
    priority_mask: npt.NDArray[np.bool_] | None = None,
) -> Seeds:
    r"""Greedily cull seeds that overlap significantly on the target sequence.

    Args:
        query_lengths (npt.NDArray[np.int32]): Query sequence lengths array.
        max_overlap_fraction (float): Maximum allowable overlap fraction. Defaults to 0.1.
        priority_mask (npt.NDArray[np.bool_] | None): Optional boolean mask for priority score boost.

    Returns:
        Seeds: Filtered, non-overlapping [`Seeds`][kaptive.core.kmers.Seeds] collection.
    """
    n = len(self)
    if n == 0:
        return self

    # Sort order: priority_mask (True first), then scores (descending)
    if priority_mask is None:
        priority_mask = np.zeros(n, dtype=np.bool_)

    order = np.lexsort((-self.scores, ~priority_mask)).astype(np.int32)

    # We pass target_indices as the 'names' array to only cull overlaps on the same contig
    intervals = self.to_intervals(query_lengths)
    kept_mask = intervals.cull_overlaps(
        order=order,
        max_overlap_fraction=max_overlap_fraction,
        group_by=self.target_indices.astype(np.int32),
    )
    return self.filter(kept_mask)

@classmethod
def empty(cls) -> Seeds:
    r"""Create an empty [`Seeds`][kaptive.core.kmers.Seeds] collection.

    Returns:
        Seeds: Empty [`Seeds`][kaptive.core.kmers.Seeds] collection.
    """
    return cls(
        np.empty(0, dtype=np.uint32),
        np.empty(0, dtype=np.uint32),
        np.empty(0, dtype=np.uint32),
        np.empty(0, dtype=np.int32),
    )

def extract_sequences(self, queries: Sequences, targets: Sequences) -> tuple[Sequences, Sequences]:
    r"""Extract parallel batches of query and target sequences mapped by this seed batch.

    Args:
        queries (Sequences): Source collection of query sequences.
        targets (Sequences): Source collection of target sequences.

    Returns:
        tuple[Sequences, Sequences]: Paired query and target sequence collections.
    """
    return queries[self.query_indices], targets[self.target_indices]  # type: ignore

def filter(self, mask: npt.NDArray[np.bool_]) -> Seeds:
    r"""Return a new Seeds collection containing only records where the mask is True.

    Args:
        mask (npt.NDArray[np.bool_]): 1D boolean mask array.

    Returns:
        Seeds: Filtered [`Seeds`][kaptive.core.kmers.Seeds] collection.
    """
    return Seeds(
        self.query_indices[mask],
        self.target_indices[mask],
        self.scores[mask],
        self.offsets[mask],
    )

def to_intervals(self, query_lengths: npt.NDArray[np.int32]) -> Intervals:
    r"""Convert seeds into target coordinates [`Intervals`][kaptive.core.interval.Intervals].

    Args:
        query_lengths (npt.NDArray[np.int32]): Query sequence lengths array.

    Returns:
        Intervals: Spatial [`Intervals`][kaptive.core.interval.Intervals] collection on target sequences.
    """
    t_starts = -self.offsets
    q_lens = query_lengths[self.query_indices]
    t_ends = t_starts + q_lens

    return Intervals(
        starts=t_starts,
        ends=t_ends,
        strands=np.ones(len(self), dtype=np.int8),
        original_indices=np.arange(len(self), dtype=np.int32),
    )

def top_hits(self, min_score: int = 1) -> Seeds:
    r"""Reduce the batch to only the highest-scoring seed for each query.

    Args:
        min_score (int): Minimum match score threshold. Defaults to 1.

    Returns:
        Seeds: Filtered [`Seeds`][kaptive.core.kmers.Seeds] collection containing top hit per query.
    """
    if len(self) == 0:
        return self

    # np.lexsort prioritizes the last array.
    # Primary key: query_indices (ascending). Secondary: scores (descending via ~)
    order = np.lexsort((~self.scores, self.query_indices))

    # Find the first occurrence of each query (which corresponds to its highest score)
    _, unique_idx = np.unique(self.query_indices[order], return_index=True)

    # Map back to the original array indices and sort to preserve the original query order
    best_idx = order[unique_idx]
    best_idx.sort()

    best_batch = Seeds(
        self.query_indices[best_idx],
        self.target_indices[best_idx],
        self.scores[best_idx],
        self.offsets[best_idx],
    )

    if min_score > 0:
        return best_batch.filter(best_batch.scores >= min_score)

    return best_batch