kaptive.serotyping.models¶

def __getitem__(self, item: Any) -> "GeneHits":
    r"""Slice or boolean-mask all parallel array fields simultaneously.

    Args:
        item (Any): Slice, integer array, or boolean mask.

    Returns:
        GeneHits: A sliced [`GeneHits`][kaptive.serotyping.models.GeneHits] instance.
    """
    return GeneHits(
        gene_indices=self.gene_indices[item],
        q_starts=self.q_starts[item],
        q_ends=self.q_ends[item],
        t_indices=self.t_indices[item],
        t_starts=self.t_starts[item],
        t_ends=self.t_ends[item],
        strands=self.strands[item],
        is_expected=self.is_expected[item],
        is_inside=self.is_inside[item],
        is_extra=self.is_extra[item],
        expected_positions=self.expected_positions[item],
        expected_strands=self.expected_strands[item],
        gene_ids=self.gene_ids[item],
        cluster_names=self.cluster_names[item],
        product_descriptions=self.product_descriptions[item],
        coverages=self.coverages[item],
    )

def __len__(self) -> int:
    r"""Return total number of gene hit alignments in container.

    Returns:
        int: Number of elements in parallel arrays.
    """
    return len(self.gene_indices)

@classmethod
def concat(cls, batches: Iterable[Self]) -> Self:  # type: ignore
    r"""Concatenate multiple `GeneHits` batches into a single container.

    Args:
        batches (Iterable[GeneHits]): An iterable of [`GeneHits`][kaptive.serotyping.models.GeneHits] instances.

    Returns:
        GeneHits: Combined [`GeneHits`][kaptive.serotyping.models.GeneHits] container.
    """
    batches_list = list(batches)
    if not batches_list:
        return cls.empty()  # type: ignore
    return cls(
        gene_indices=np.concatenate([b.gene_indices for b in batches_list]),
        q_starts=np.concatenate([b.q_starts for b in batches_list]),
        q_ends=np.concatenate([b.q_ends for b in batches_list]),
        t_indices=np.concatenate([b.t_indices for b in batches_list]),
        t_starts=np.concatenate([b.t_starts for b in batches_list]),
        t_ends=np.concatenate([b.t_ends for b in batches_list]),
        strands=np.concatenate([b.strands for b in batches_list]),
        is_expected=np.concatenate([b.is_expected for b in batches_list]),
        is_inside=np.concatenate([b.is_inside for b in batches_list]),
        is_extra=np.concatenate([b.is_extra for b in batches_list]),
        expected_positions=np.concatenate([b.expected_positions for b in batches_list]),
        expected_strands=np.concatenate([b.expected_strands for b in batches_list]),
        gene_ids=np.concatenate([b.gene_ids for b in batches_list]),
        cluster_names=np.concatenate([b.cluster_names for b in batches_list]),
        product_descriptions=np.concatenate([b.product_descriptions for b in batches_list]),
        coverages=np.concatenate([b.coverages for b in batches_list]),
    )

@classmethod
def empty(cls) -> "GeneHits":
    r"""Create an empty `GeneHits` container with zero-length arrays and empty tuples.

    Returns:
        GeneHits: An empty [`GeneHits`][kaptive.serotyping.models.GeneHits] instance.
    """
    return cls(
        np.empty(0, dtype=np.int32),
        np.empty(0, dtype=np.int32),
        np.empty(0, dtype=np.int32),
        np.empty(0, dtype=np.uint32),
        np.empty(0, dtype=np.int32),
        np.empty(0, dtype=np.int32),
        np.empty(0, dtype=np.int8),
        np.empty(0, dtype=bool),
        np.empty(0, dtype=bool),
        np.empty(0, dtype=bool),
        np.empty(0, dtype=np.int32),
        np.empty(0, dtype=np.int8),
        np.empty(0, dtype="S32"),
        np.empty(0, dtype="S10"),
        np.empty(0, dtype="S64"),
        np.empty(0, dtype=np.float32),
    )

@classmethod
def from_dict(cls, data: dict[str, Any]) -> "GeneHits":
    r"""Reconstruct a `GeneHits` container from a deserialized dictionary.

    Args:
        data (dict[str, Any]): Dictionary containing array data lists and tuple metadata.

    Returns:
        GeneHits: Reconstructed [`GeneHits`][kaptive.serotyping.models.GeneHits] instance.
    """

    def _to_bytes_array(val: Any, dtype: str) -> npt.NDArray[np.bytes_]:
        if val is None or len(val) == 0:
            return np.empty(0, dtype=dtype)
        if isinstance(val, np.ndarray) and val.dtype.kind in ("S", "a"):
            return val.astype(dtype)
        encoded = [x.encode("utf-8") if isinstance(x, str) else x for x in val]
        return np.array(encoded, dtype=dtype)

    return cls(
        gene_indices=np.array(data["gene_indices"], dtype=np.int32),
        q_starts=np.array(data["q_starts"], dtype=np.int32),
        q_ends=np.array(data["q_ends"], dtype=np.int32),
        t_indices=np.array(data["t_indices"], dtype=np.uint32),
        t_starts=np.array(data["t_starts"], dtype=np.int32),
        t_ends=np.array(data["t_ends"], dtype=np.int32),
        strands=np.array(data["strands"], dtype=np.int8),
        is_expected=np.array(data["is_expected"], dtype=bool),
        is_inside=np.array(data["is_inside"], dtype=bool),
        is_extra=np.array(data["is_extra"], dtype=bool),
        expected_positions=np.array(data.get("expected_positions", []), dtype=np.int32),
        expected_strands=np.array(data.get("expected_strands", []), dtype=np.int8),
        gene_ids=_to_bytes_array(data.get("gene_ids", []), "S32"),
        cluster_names=_to_bytes_array(data.get("cluster_names", []), "S10"),
        product_descriptions=_to_bytes_array(data.get("product_descriptions", []), "S64"),
        coverages=np.array(data.get("coverages", []), dtype=np.float32),
    )

def to_dict(self) -> dict[str, Any]:
    r"""Convert SoA array fields to a dictionary for JSON serialization.

    Returns:
        dict[str, Any]: Dictionary mapping field names to NumPy arrays and metadata tuples.
    """
    d = {
        k: getattr(self, k)
        for k in (
            "gene_indices",
            "q_starts",
            "q_ends",
            "t_indices",
            "t_starts",
            "t_ends",
            "strands",
            "is_expected",
            "is_inside",
            "is_extra",
            "expected_positions",
            "expected_strands",
            "coverages",
        )
    }
    d["gene_ids"] = np.char.decode(self.gene_ids, "utf-8").tolist()
    d["cluster_names"] = np.char.decode(self.cluster_names, "utf-8").tolist()
    d["product_descriptions"] = np.char.decode(self.product_descriptions, "utf-8").tolist()
    return d

def __getitem__(self, item: int | slice | npt.NDArray[Any] | list[int]) -> "Any | LocusPieces":
    r"""Slice or array-mask all parallel fields of locus pieces simultaneously.

    Args:
        item (int | slice | npt.NDArray | list): Slice range, boolean mask, or index list.

    Returns:
        LocusPieces: A sliced [`LocusPieces`][kaptive.serotyping.models.LocusPieces] instance.

    Raises:
        NotImplementedError: If single integer key access is attempted.
    """
    if isinstance(item, (int, np.integer)):
        raise NotImplementedError("Single item access not implemented for LocusPieces")
    return LocusPieces(
        ctg_indices=self.ctg_indices[item],
        starts=self.starts[item],
        ends=self.ends[item],
        strands=self.strands[item],
    )

def __len__(self) -> int:
    r"""Return total number of locus pieces in container.

    Returns:
        int: Number of fragment elements.
    """
    return len(self.ctg_indices)

@classmethod
def concat(cls, batches: Iterable[Self]) -> Self:  # type: ignore
    r"""Concatenate multiple `LocusPieces` batches into a single container.

    Args:
        batches (Iterable[LocusPieces]): An iterable of
            [`LocusPieces`][kaptive.serotyping.models.LocusPieces] instances.

    Returns:
        LocusPieces: Combined [`LocusPieces`][kaptive.serotyping.models.LocusPieces] container.
    """
    batches_list = list(batches)
    if not batches_list:
        return cls.empty()  # type: ignore
    return cls(
        ctg_indices=np.concatenate([b.ctg_indices for b in batches_list]),
        starts=np.concatenate([b.starts for b in batches_list]),
        ends=np.concatenate([b.ends for b in batches_list]),
        strands=np.concatenate([b.strands for b in batches_list]),
    )

@classmethod
def empty(cls) -> "LocusPieces":
    r"""Create an empty `LocusPieces` container with zero-length arrays.

    Returns:
        LocusPieces: An empty [`LocusPieces`][kaptive.serotyping.models.LocusPieces] instance.
    """
    return cls(
        np.empty(0, dtype=np.uint32),
        np.empty(0, dtype=np.int32),
        np.empty(0, dtype=np.int32),
        np.empty(0, dtype=np.int8),
    )

@classmethod
def from_dict(cls, data: dict[str, Any]) -> "LocusPieces":
    r"""Reconstruct a `LocusPieces` container from a deserialized dictionary.

    Args:
        data (dict[str, Any]): Dictionary containing coordinate array lists.

    Returns:
        LocusPieces: Reconstructed [`LocusPieces`][kaptive.serotyping.models.LocusPieces] instance.
    """
    return cls(
        ctg_indices=np.array(data["ctg_indices"], dtype=np.uint32),
        starts=np.array(data["starts"], dtype=np.int32),
        ends=np.array(data["ends"], dtype=np.int32),
        strands=np.array(data["strands"], dtype=np.int8),
    )

def to_dict(self) -> dict[str, Any]:
    r"""Convert array fields to a dictionary for JSON serialization.

    Returns:
        dict[str, Any]: Dictionary mapping field names to NumPy array data.
    """
    return {k: getattr(self, k) for k in ("ctg_indices", "starts", "ends", "strands")}

def to_symbols(self) -> bytes:
    r"""Render the bitflag combination into formatted symbol bytes for TSV reporting.

    Returns:
        bytes: ASCII byte string containing concatenation of active problem symbols.
    """
    return self.SYMBOLS[self.value]

@classmethod
def from_dict(cls, data: dict[str, Any]) -> "SerotypingResult":
    r"""Reconstruct a `SerotypingResult` instance from a deserialized dictionary.

    Args:
        data (dict[str, Any]): Dictionary containing serialized fields and nested sub-dictionaries.

    Returns:
        SerotypingResult: Reconstructed [`SerotypingResult`][kaptive.serotyping.models.SerotypingResult] instance.
    """
    return cls(
        kaptive_version=data["kaptive_version"],
        database_name=data["database_name"],
        database_version=data["database_version"],
        database_organism=data["database_organism"],
        database_taxon=data["database_taxon"],
        genome=data["genome"],
        best_locus_idx=data["best_locus_idx"],
        best_locus_name=data["best_locus_name"],
        best_locus_score=data["best_locus_score"],
        best_locus_completeness=data["best_locus_completeness"],
        length_discrepancy=data["length_discrepancy"],
        locus_pieces=LocusPieces.from_dict(data["locus_pieces"]),
        gene_hits=GeneHits.from_dict(data["gene_hits"]),
        gene_states=np.array(data["gene_states"], dtype=np.int8),
        percent_identity=data["percent_identity"],
        percent_coverage=data["percent_coverage"],
        phenotype=data["phenotype"],
        typeable=data["typeable"],
        missing_expected_genes=tuple(data.get("missing_expected_genes", [])),
        locus_seqs=Sequences.from_dict(data["locus_seqs"]),
        gene_seqs=Sequences.from_dict(data["gene_seqs"]),
        translations=Sequences.from_dict(data["translations"]),
        protein_identities=np.array(data["protein_identities"], dtype=np.float32),
    )

def to_dict(self) -> dict[str, Any]:
    r"""Convert serotyping result into a dictionary suitable for JSON serialization.

    Returns:
        dict[str, Any]: Lightweight dictionary containing primitive types, lists, and nested dictionaries.
    """
    return {
        "kaptive_version": self.kaptive_version,
        "database_name": self.database_name,
        "database_version": self.database_version,
        "database_organism": self.database_organism,
        "database_taxon": self.database_taxon,
        "genome": self.genome,
        "best_locus_idx": self.best_locus_idx,
        "best_locus_name": self.best_locus_name,
        "best_locus_score": self.best_locus_score,
        "best_locus_completeness": self.best_locus_completeness,
        "length_discrepancy": self.length_discrepancy,
        "percent_identity": self.percent_identity,
        "percent_coverage": self.percent_coverage,
        "phenotype": self.phenotype,
        "typeable": self.typeable,
        "missing_expected_genes": self.missing_expected_genes,
        "problems": self.problems,
        "locus_pieces": self.locus_pieces.to_dict(),
        "gene_hits": self.gene_hits.to_dict(),
        "gene_states": self.gene_states,
        "protein_identities": self.protein_identities,
        "locus_seqs": self.locus_seqs.to_dict(),
        "gene_seqs": self.gene_seqs.to_dict(),
        "translations": self.translations.to_dict(),
    }

def to_locus_data(self) -> "LocusData":
    r"""Convert result into a `LocusData` container for comparative multi-locus visualization.

    Extracts translations, locus backbone intervals, locus pieces, contig indices, gene states,
    and functional product descriptions for non-extra inside genes.

    Returns:
        LocusData: A [`LocusData`][kaptive.compare.LocusData] instance for multi-locus alignment plotting.
    """
    from kaptive.compare import LocusData

    mask = self.gene_hits.is_inside & ~self.gene_hits.is_extra
    descriptions = np.asarray(
        np.char.decode(self.gene_hits.product_descriptions[mask], "utf-8"),
        dtype=object,
    )

    return LocusData(
        proteins=self.translations[mask],  # type: ignore
        name=self.genome,
        backbone=self.gene_hits.t_intervals[mask],  # type: ignore
        pieces=self.locus_pieces,
        gene_ctg_indices=self.gene_hits.t_indices[mask],
        gene_states=self.gene_states[mask],
        gene_descriptions=descriptions,
    )