kaptive.serotyping¶

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 __bytes__(self) -> bytes:
    r"""Serialize the report row fields into a tab-separated binary TSV row.

    Returns:
        bytes: Tab-separated field values ending with a newline (`b"\n"`).
    """
    return b"\t".join(getattr(self, f.name) for f in fields(self)) + b"\n"

@classmethod
def from_result(cls, result: SerotypingResult) -> "KaptiveRow":
    r"""Construct a classic [`KaptiveRow`][kaptive.serotyping.io.KaptiveRow] from a serotyping result.

    Calculates gene counts, percentage coverages, identity metrics, and problem symbol codes, formatting all fields
    into UTF-8 encoded bytes for backwards-compatible TSV output.

    Args:
        result (SerotypingResult): The serotyping call result.
            See [`SerotypingResult`][kaptive.serotyping.models.SerotypingResult].

    Returns:
        KaptiveRow: Formatted report row object.
    """
    hits = result.gene_hits
    states = result.gene_states

    # High-performance boolean masks
    in_loc = hits.is_inside
    out_loc = ~hits.is_inside
    exp = hits.is_expected
    extra = hits.is_extra
    unexp = ~exp & ~extra

    def _format_genes(mask: np.ndarray) -> bytes:
        r"""Helper to rapidly construct the details string for a specific masked subset of genes.

        Args:
            mask (np.ndarray): Boolean mask selecting the subset of gene hits to format.

        Returns:
            bytes: Semicolon-separated details string for the selected genes.
        """
        indices = np.where(mask)[0]
        if indices.size == 0:
            return b""

        details = []
        for i in indices:
            gene_name = result.gene_seqs.ids[i].encode("utf-8")
            parts = [
                gene_name,
                b"%.2f%%" % result.protein_identities[i],
                b"%.2f%%" % result.gene_hits.coverages[i],
            ]

            if states[i] == GeneState.PARTIAL.value:
                parts.append(b"partial")
            elif states[i] == GeneState.TRUNCATED.value:
                parts.append(b"truncated")
            elif states[i] == GeneState.NOVEL.value:
                parts.append(b"below_id_threshold")

            details.append(b",".join(parts))
        return b";".join(details)

    # Expected Inside
    mask_exp_in = in_loc & exp
    n_exp_in = len(np.unique(result.gene_hits.gene_indices[mask_exp_in]))

    # Expected Outside
    mask_exp_out = out_loc & exp
    n_exp_out = len(np.unique(result.gene_hits.gene_indices[mask_exp_out]))

    expected_total = n_exp_in + n_exp_out + len(result.missing_expected_genes)

    in_comp = (n_exp_in / expected_total * 100.0) if expected_total > 0 else 0.0
    exp_in_str = b"%d / %d (%.2f%%)" % (n_exp_in, expected_total, in_comp) if expected_total else b"0 / 0 (0.00%)"

    out_comp = (n_exp_out / expected_total * 100.0) if expected_total > 0 else 0.0
    exp_out_str = (
        b"%d / %d (%.2f%%)" % (n_exp_out, expected_total, out_comp) if expected_total else b"0 / 0 (0.00%)"
    )

    # Other counts
    n_unexp_in = len(np.unique(result.gene_hits.gene_indices[in_loc & unexp]))
    n_unexp_out = len(np.unique(result.gene_hits.gene_indices[out_loc & unexp]))

    return cls(
        Kaptive_version=result.kaptive_version.encode(),
        Database_name=result.database_name.encode(),
        Database_version=result.database_version.encode(),
        Assembly=result.genome.encode(),
        Best_match_locus=result.best_locus_name.encode(),
        Best_match_type=result.phenotype.encode(),
        Match_confidence=b"Typeable" if result.typeable else b"Untypeable",
        Problems=result.problems.to_symbols(),
        Identity=b"%.2f%%" % result.percent_identity,
        Coverage=b"%.2f%%" % result.percent_coverage,
        Length_discrepancy=b"n/a"
        if (result.length_discrepancy is None or np.isnan(result.length_discrepancy))
        else b"%d" % int(result.length_discrepancy),
        Expected_genes_in_locus=exp_in_str,
        Expected_genes_in_locus_details=_format_genes(mask_exp_in),
        Missing_expected_genes=b";".join(g.encode("utf-8") for g in result.missing_expected_genes),
        Other_genes_in_locus=b"%d" % n_unexp_in,
        Other_genes_in_locus_details=_format_genes(in_loc & unexp),
        Expected_genes_outside_locus=exp_out_str,
        Expected_genes_outside_locus_details=_format_genes(mask_exp_out),
        Other_genes_outside_locus=b"%d" % n_unexp_out,
        Other_genes_outside_locus_details=_format_genes(out_loc & unexp),
        Truncated_genes_details=_format_genes(
            (states == GeneState.TRUNCATED.value) | (states == GeneState.PARTIAL.value)
        ),
        Extra_genes_details=_format_genes(extra),
    )

@classmethod
def header(cls) -> bytes:
    r"""Generate backwards-compatible column header bytes for classic Kaptive reports.

    Replaces internal field name underscores with spaces and `_details` with `, details` to maintain exact
    compatibility with legacy Kaptive TSV headers.

    Returns:
        bytes: Tab-separated legacy header line ending with a newline (`b"\n"`).
    """
    headers = [f.name.encode("utf-8").replace(b"_details", b", details").replace(b"_", b" ") for f in fields(cls)]
    return b"\t".join(headers) + b"\n"

@classmethod
def read_tsv(cls, lines: Iterable[bytes]) -> Iterator[Self]:
    r"""Parse an iterable of TSV byte lines into report row instances.

    Args:
        lines (Iterable[bytes]): An iterable of byte strings (e.g. an open binary file).

    Yields:
        Self: Parsed instances of the concrete [`ReportRow`][kaptive.serotyping.io.ReportRow] subclass.
    """
    line_iter = iter(lines)
    try:
        header_line = next(line_iter)
    except StopIteration:
        return

    field_names = cls._parse_header_line(header_line)
    expected_fields = {f.name for f in fields(cls)}

    for line in line_iter:
        line = line.rstrip(b"\r\n")
        if not line:
            continue
        values = line.split(b"\t")
        kwargs = {name: val for name, val in zip(field_names, values) if name in expected_fields}
        yield cls(**kwargs)

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 __bytes__(self) -> bytes:
    r"""Serialize the report row fields into a tab-separated binary TSV row.

    Returns:
        bytes: Tab-separated field values ending with a newline (`b"\n"`).
    """
    return b"\t".join(getattr(self, f.name) for f in fields(self)) + b"\n"

@classmethod
def from_result(cls, result: SerotypingResult) -> "Pha4geRow":
    r"""Construct a standardized [`Pha4geRow`][kaptive.serotyping.io.Pha4geRow] from a serotyping result.

    Transforms numeric taxon IDs and problem flags into human-readable PHA4GE-compliant strings and binary bytes.

    Args:
        result (SerotypingResult): The serotyping call result.
            See [`SerotypingResult`][kaptive.serotyping.models.SerotypingResult].

    Returns:
        Pha4geRow: Formatted PHA4GE report row object.
    """
    # Transform problem enum into human-readable string
    if result.problems:
        detail_parts = []
        if SerotypingProblem.TRUNCATED_GENES in result.problems:
            detail_parts.append(b"truncated gene/s in locus")
        if SerotypingProblem.NOVEL_GENES in result.problems:
            detail_parts.append(b"low identity gene/s")
        if SerotypingProblem.FRAGMENTED in result.problems:
            detail_parts.append(b"match broken into %d pieces" % len(result.locus_pieces))
        if SerotypingProblem.MISSING_GENES in result.problems:
            detail_parts.append(b"missing expected gene/s")
        if SerotypingProblem.UNEXPECTED_GENES in result.problems:
            detail_parts.append(b"unexpected gene/s in locus")
        details = b"Best locus match: %b. Problems: %b" % (
            result.best_locus_name.encode(),
            b", ".join(detail_parts),
        )
    else:
        details = b"Best locus match: %b." % result.best_locus_name.encode()

    return cls(
        sample=result.genome.encode(),
        genotyping_schema_taxon=b"%s [NCBITaxon:%d]" % (result.database_organism.encode(), result.database_taxon),
        genotyping_database_name=result.database_name.encode(),
        genotyping_database_version=result.database_version.encode(),
        genotyping_software_version=result.kaptive_version.encode(),
        genotype=result.best_locus_name.encode(),
        genotype_confidence_value=b"Typeable" if result.typeable else b"Untypeable",
        genotype_predicted_phenotype=result.phenotype.encode(),
        genotyping_details=details,
    )

@classmethod
def header(cls) -> bytes:
    r"""Generate the TSV header row as UTF-8 encoded bytes.

    Returns:
        bytes: Tab-separated column header line ending with a newline (`b"\n"`).
    """
    return ("\t".join(f.name for f in fields(cls)) + "\n").encode("utf-8")

@classmethod
def read_tsv(cls, lines: Iterable[bytes]) -> Iterator[Self]:
    r"""Parse an iterable of TSV byte lines into report row instances.

    Args:
        lines (Iterable[bytes]): An iterable of byte strings (e.g. an open binary file).

    Yields:
        Self: Parsed instances of the concrete [`ReportRow`][kaptive.serotyping.io.ReportRow] subclass.
    """
    line_iter = iter(lines)
    try:
        header_line = next(line_iter)
    except StopIteration:
        return

    field_names = cls._parse_header_line(header_line)
    expected_fields = {f.name for f in fields(cls)}

    for line in line_iter:
        line = line.rstrip(b"\r\n")
        if not line:
            continue
        values = line.split(b"\t")
        kwargs = {name: val for name, val in zip(field_names, values) if name in expected_fields}
        yield cls(**kwargs)

def __bytes__(self) -> bytes:
    r"""Serialize the report row fields into a tab-separated binary TSV row.

    Returns:
        bytes: Tab-separated field values ending with a newline (`b"\n"`).
    """
    return b"\t".join(getattr(self, f.name) for f in fields(self)) + b"\n"

@classmethod
@abstractmethod
def from_result(cls, result: SerotypingResult) -> Self:
    r"""Construct a report row instance from a serotyping result.

    Args:
        result (SerotypingResult): The serotyping analysis result to format.
            See [`SerotypingResult`][kaptive.serotyping.models.SerotypingResult].

    Returns:
        Self: Instance of the concrete [`ReportRow`][kaptive.serotyping.io.ReportRow] subclass
            populated with result fields.
    """
    ...

@classmethod
def header(cls) -> bytes:
    r"""Generate the TSV header row as UTF-8 encoded bytes.

    Returns:
        bytes: Tab-separated column header line ending with a newline (`b"\n"`).
    """
    return ("\t".join(f.name for f in fields(cls)) + "\n").encode("utf-8")

@classmethod
def read_tsv(cls, lines: Iterable[bytes]) -> Iterator[Self]:
    r"""Parse an iterable of TSV byte lines into report row instances.

    Args:
        lines (Iterable[bytes]): An iterable of byte strings (e.g. an open binary file).

    Yields:
        Self: Parsed instances of the concrete [`ReportRow`][kaptive.serotyping.io.ReportRow] subclass.
    """
    line_iter = iter(lines)
    try:
        header_line = next(line_iter)
    except StopIteration:
        return

    field_names = cls._parse_header_line(header_line)
    expected_fields = {f.name for f in fields(cls)}

    for line in line_iter:
        line = line.rstrip(b"\r\n")
        if not line:
            continue
        values = line.split(b"\t")
        kwargs = {name: val for name, val in zip(field_names, values) if name in expected_fields}
        yield cls(**kwargs)