kaptive.core.seq¶

@classmethod
def is_coding(cls, seq: bytes) -> bool:
    r"""Check if a byte sequence starts with a valid bacterial start codon and ends with a stop codon.

    Args:
        seq (bytes): Nucleotide sequence to check.

    Returns:
        bool: True if sequence length >= 3 and starts/ends with valid bacterial start/stop codons.
    """
    if len(seq) < 3:
        return False
    return seq[:3] in cls._START_CODONS and seq[-3:] in cls._STOP_CODONS

@classmethod
def translate(
    cls, seq: bytes | bytearray | memoryview | npt.NDArray[np.uint8], to_stop: bool = False
) -> npt.NDArray[np.uint8]:
    r"""Translate a nucleotide sequence array or byte string to amino acid byte values.

    Args:
        seq (bytes | bytearray | memoryview | npt.NDArray[np.uint8]): Input nucleotide sequence.
        to_stop (bool): If True, truncates translation at the first stop codon ('*'). Defaults to False.

    Returns:
        npt.NDArray[np.uint8]: 1D array of ASCII uint8 values representing translated amino acids.
    """
    if len(seq) < 3:
        return np.array([], dtype=np.uint8)

    if not isinstance(seq, np.ndarray):
        seq = np.ascontiguousarray(np.frombuffer(seq, np.uint8))

    return _translate_kernel(seq, cls._CHAR_MAP, cls._CODON_MAP, to_stop)

def __len__(self) -> int:
    r"""Return the length of the sequence in bytes.

    Returns:
        int: Sequence length in bytes.
    """
    return len(self.seq)

def extract(self, start: int | IntervalLike, end: int | None = None, strand: Strand = Strand.UNSTRANDED) -> bytes:
    r"""Extract a sub-sequence based on coordinates and orientation.

    If `strand` is negative (e.g. [`Strand`][kaptive.core.interval.Strand]),
    the extracted sub-sequence is automatically reverse-complemented.

    Args:
        start (int | IntervalLike): 0-based start coordinate, or an
            `IntervalLike` object (providing start, end, and strand).
        end (int | None): 0-based end coordinate (exclusive). If `start` is an interval,
            this must be None. Defaults to None.
        strand (Strand): Orientation of the extraction. Defaults to
            [`Strand`][kaptive.core.interval.Strand].

    Returns:
        bytes: Extracted (and potentially reverse-complemented) sub-sequence bytes.
    """
    if end is None:
        interval = Interval.from_item(start, strand=strand)
        start_val, end_val, strand_val = interval.start, interval.end, interval.strand
    else:
        start_val, end_val, strand_val = int(start), int(end), strand  # type: ignore

    new_seq = self.seq[start_val:end_val]
    if strand_val < 0:
        return bytes(new_seq.translate(BacterialTranslationTable._COMP)[::-1])
    return bytes(new_seq)

def to_fasta(self) -> bytes:
    r"""Format the sequence as a FASTA record byte string.

    Returns:
        bytes: A byte string containing the FASTA header (`>id\n`) followed by
            the sequence and a trailing newline.
    """
    return b">%b\n%b\n" % (self.id.encode(), self.seq)

def __getitem__(self, item: int | slice | np.ndarray[Any, Any] | list[int]) -> SeqRecord | Sequences:
    r"""Access sequences by index, slice, or boolean mask.

    Args:
        item (int | slice | np.ndarray | list): An integer index, slice, or NumPy array mask/indices.

    Returns:
        SeqRecord | Sequences: A single [`SeqRecord`][kaptive.core.seq.SeqRecord] if integer index,
            or a sliced [`Sequences`][kaptive.core.seq.Sequences] collection.

    Raises:
        IndexError: If an integer index is out of bounds.
    """
    if isinstance(item, (int, np.integer)):
        item_idx = int(item)
        if item_idx < 0:
            item_idx += len(self)
        if item_idx < 0 or item_idx >= len(self):
            raise IndexError("Batch index out of range")
        offset_val = self.offsets[item_idx]
        length_val = self.lengths[item_idx]
        return SeqRecord(self.ids[item_idx], self.seqs[offset_val : offset_val + length_val].tobytes())

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

    starts = np.zeros(len(indices), dtype=np.int32)
    ends = self.lengths[indices].astype(np.int32)
    strands = np.ones(len(indices), dtype=np.int8)
    return self.extract(
        indices.astype(np.int32), starts, ends, strands, new_ids=tuple(self.ids[i] for i in indices)
    )

def __iter__(self) -> Generator[SeqRecord, None, None]:
    r"""Iterate over the batch, yielding a SeqRecord for each sequence.

    Yields:
        SeqRecord: Scalar record for each sequence in the batch.
    """
    for i in range(len(self)):
        offset_val = self.offsets[i]
        length_val = self.lengths[i]
        yield SeqRecord(self.ids[i], self.seqs[offset_val : offset_val + length_val].tobytes())

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

    Returns:
        int: Number of sequence records in the container.
    """
    return len(self.ids)

@classmethod
def concat(cls, batches: Iterable[Self]) -> Sequences:  # type: ignore
    r"""Concatenate multiple Sequences containers into a single larger collection.

    Args:
        batches (Iterable[Sequences]): An iterable of sequence batches to combine.

    Returns:
        Sequences: A combined sequences container.
    """
    batches_list = list(batches)
    if not batches_list:
        return cls.empty()

    all_ids = sum((b.ids for b in batches_list), ())
    all_seqs = np.concatenate([b.seqs for b in batches_list])
    all_lengths = np.concatenate([b.lengths for b in batches_list])

    offsets = np.zeros(len(all_lengths), dtype=np.int32)
    if len(all_lengths) > 1:
        np.cumsum(all_lengths[:-1], out=offsets[1:])

    return cls(all_ids, all_seqs, offsets, all_lengths)

@classmethod
def empty(cls) -> Sequences:
    r"""Create an empty Sequences object with zero-length arrays.

    Returns:
        Sequences: An empty sequences container.
    """
    return cls(
        (),
        np.empty(0, dtype=np.uint8),
        np.empty(0, dtype=np.int32),
        np.empty(0, dtype=np.int32),
    )

def extract(
    self,
    indices: npt.NDArray[np.int32],
    starts: npt.NDArray[np.int32],
    ends: npt.NDArray[np.int32],
    strands: npt.NDArray[np.int8],
    new_ids: tuple[str, ...] | None = None,
) -> Sequences:
    r"""Vectorized sub-sequence extraction from the batch.

    Handles coordinates and reverse-complementation across multiple sequences in parallel using Numba.

    Args:
        indices (npt.NDArray[np.int32]): 1D array specifying parent sequence index for each extraction.
        starts (npt.NDArray[np.int32]): 0-based start coordinates relative to parent sequence.
        ends (npt.NDArray[np.int32]): 0-based end coordinates relative to parent sequence.
        strands (npt.NDArray[np.int8]): Strand orientations (1 for forward, -1 for reverse-complement).
        new_ids (tuple[str, ...] | None): Identifiers for extracted sequences. If None, auto-generates names.

    Returns:
        Sequences: New collection of extracted sub-sequences.
    """
    if len(indices) == 0:
        return self.empty()
    new_ids = new_ids or tuple(f"{self.ids[i]}_{x}_{y}_{z}" for i, x, y, z in zip(indices, starts, ends, strands))
    out_seqs, offsets, lengths = _extract_ragged_kernel(
        self.seqs, self.offsets, indices, starts, ends, strands, BacterialTranslationTable._COMP_MAP
    )
    return Sequences(new_ids, out_seqs, offsets, lengths)

def extract_intervals(
    self,
    indices: npt.NDArray[np.integer],
    intervals: Intervals,
    new_ids: tuple[str, ...] | None = None,
) -> Sequences:
    r"""Wrapper around extract taking an Intervals collection.

    Args:
        indices (npt.NDArray[np.integer]): Target sequence indices for each interval.
        intervals (Intervals): Interval collection containing coordinates and strands.
        new_ids (tuple[str, ...] | None): New sequence identifiers. Defaults to None.

    Returns:
        Sequences: New collection of extracted sub-sequences.
    """
    return self.extract(
        indices.astype(np.int32),
        intervals.starts.astype(np.int32),
        intervals.ends.astype(np.int32),
        intervals.strands,
        new_ids=new_ids,
    )

@classmethod
def from_bytes(cls, seqs: list[bytes], ids: tuple[str, ...] | None = None) -> Sequences:
    r"""Construct Sequences from a list of byte strings.

    Args:
        seqs (list[bytes]): List of raw sequence byte strings.
        ids (tuple[str, ...] | None): Sequence identifiers. Defaults to string integer indices ("0", "1", ...).

    Returns:
        Sequences: Newly constructed sequences container.
    """
    ids = ids or tuple(str(i) for i in range(len(seqs)))
    return cls.from_records([SeqRecord(i, s) for i, s in zip(ids, seqs)])

@classmethod
def from_dict(cls, data: dict[str, Any]) -> Sequences:
    r"""Deserialize a Sequences object from a dictionary representation.

    Args:
        data (dict[str, Any]): Dictionary containing 'ids', ASCII string 'seqs', 'offsets', and 'lengths'.

    Returns:
        Sequences: Deserialized sequence collection.
    """
    return cls(
        ids=tuple(data["ids"]),
        seqs=np.frombuffer(data["seqs"].encode("ascii"), dtype=np.uint8),
        offsets=np.array(data["offsets"], dtype=np.int32),
        lengths=np.array(data["lengths"], dtype=np.int32),
    )

@classmethod
def from_records(cls, records: list[SeqRecord]) -> Sequences:
    r"""Construct Sequences from a list of SeqRecord objects.

    Args:
        records (list[SeqRecord]): List of sequence records.

    Returns:
        Sequences: Newly constructed sequences container.
    """
    ids = tuple(r.id for r in records)
    seqs = [np.frombuffer(r.seq, dtype=np.uint8) for r in records]
    if not seqs:
        return cls.empty()
    out_seqs = np.concatenate(seqs, dtype=np.uint8)
    lengths = np.array([len(s) for s in seqs], dtype=np.int32)
    offsets = np.zeros(len(seqs), dtype=np.int32)
    if len(seqs) > 1:
        np.cumsum(lengths[:-1], out=offsets[1:])
    return cls(ids, out_seqs, offsets, lengths)

def to_dict(self) -> dict[str, Any]:
    r"""Convert sequence batch to a dictionary representation suitable for serialization.

    Returns:
        dict[str, Any]: Dictionary containing 'ids', ASCII string 'seqs', 'offsets', and 'lengths'.
    """
    return {
        "ids": self.ids,
        "seqs": self.seqs.tobytes().decode("ascii"),
        "offsets": self.offsets,
        "lengths": self.lengths,
    }

def to_fasta(self, use_indices: bool = False) -> bytes:
    r"""Format the entire batch as a single FASTA byte string.

    Args:
        use_indices (bool): If True, uses 0-based integer index as FASTA header (`>0`, `>1`, ...)
            instead of string `ids`. Defaults to False.

    Returns:
        bytes: Complete FASTA-formatted byte string containing all sequences.
    """
    if not self.ids and not use_indices:
        return b""

    seq_bytes = self.seqs.tobytes()
    if use_indices:
        return b"".join(
            [
                b">%d\n%b\n" % (i, seq_bytes[o : o + length_val])
                for i, (o, length_val) in enumerate(zip(self.offsets.tolist(), self.lengths.tolist()))
            ]
        )
    else:
        return b"".join(
            [
                b">%b\n%b\n" % (i.encode(), seq_bytes[o : o + length_val])
                for i, o, length_val in zip(self.ids, self.offsets.tolist(), self.lengths.tolist())
            ]
        )

def translate(self, frames: npt.NDArray[np.int8] | None = None, to_stop: bool = False) -> Sequences:
    r"""Vectorized translation of nucleotide sequences into protein sequences.

    Translates sequences according to NCBI Translation Table 11 (Bacterial, Archaeal, and Plant Plastid Code)
    via [`BacterialTranslationTable`][kaptive.core.seq.BacterialTranslationTable].

    Args:
        frames (npt.NDArray[np.int8] | None): Reading frame offsets (0, 1, or 2) per sequence.
            Defaults to frame 0 for all.
        to_stop (bool): If True, truncates translation at the first stop codon (*). Defaults to False.

    Returns:
        Sequences: New collection containing translated protein sequences.
    """
    if len(self) == 0:
        return self.empty()
    if frames is None:
        frames = np.zeros(len(self), dtype=np.int8)
    out_seqs, offsets, lengths = _translate_ragged_kernel(
        self.seqs,
        self.offsets,
        self.lengths,
        frames,
        BacterialTranslationTable._CHAR_MAP,
        BacterialTranslationTable._CODON_MAP,
        to_stop,
    )
    return Sequences(self.ids, out_seqs, offsets, lengths)

def unique(self) -> Sequences:
    r"""Return a new Sequences container containing only unique sequences.

    Uses a Numba-accelerated 64-bit FNV-1a hash to identify unique sequences while preserving
    first-occurrence order.

    Returns:
        Sequences: A new container with duplicate sequences removed.
    """
    if len(self) <= 1:
        return self

    hashes = _hash_sequences_kernel(self.seqs, self.offsets, self.lengths)
    _, unique_indices = np.unique(hashes, return_index=True)
    unique_indices.sort()

    return self[unique_indices]  # type: ignore