Use enums for fixed pipeline choices

copykat_py/__init__.py

@@ -1,9 +1,44 @@
 """CopyKAT-Py: Python implementation of CopyKAT for scRNA-seq copy number inference."""
 
+from copykat_py._types import (
+    AnchorPath,
+    AnchorStrategy,
+    BaselineWarning,
+    CellLineMode,
+    DataQualityStatus,
+    DistanceMetric,
+    ExecutionBackend,
+    FinalCallStrategy,
+    GeneIdType,
+    Genome,
+    KSMethod,
+    PredictionLabel,
+    ReferenceMode,
+    ReportFormat,
+)
 from copykat_py.cli import copykat_anndata
 from copykat_py.copykat import copykat
 from copykat_py.data_loader import load_example_data
 from copykat_py.plotting import plot_heatmap_annotated
 
 __version__ = "1.0.0"
-__all__ = ["copykat", "copykat_anndata", "load_example_data", "plot_heatmap_annotated"]
+__all__ = [
+    "AnchorPath",
+    "AnchorStrategy",
+    "BaselineWarning",
+    "CellLineMode",
+    "DataQualityStatus",
+    "DistanceMetric",
+    "ExecutionBackend",
+    "FinalCallStrategy",
+    "GeneIdType",
+    "Genome",
+    "KSMethod",
+    "PredictionLabel",
+    "ReferenceMode",
+    "ReportFormat",
+    "copykat",
+    "copykat_anndata",
+    "load_example_data",
+    "plot_heatmap_annotated",
+]

copykat_py/_types.py

@@ -8,8 +8,11 @@
 - Clustering works on the transpose, ``CellByFeature`` (cells x features).
 """
 
+from __future__ import annotations
+
 from collections.abc import Sequence
-from typing import Any, Literal, NotRequired, TypeAlias, TypedDict
+from enum import StrEnum
+from typing import Any, NotRequired, TypeAlias, TypedDict
 
 import numpy as np
 import numpy.typing as npt
@@ -43,6 +46,124 @@
 SparseMatrix: TypeAlias = sparse.coo_matrix | sparse.csr_matrix | sparse.csc_matrix
 
 
+class DistanceMetric(StrEnum):
+    """Supported distance metrics for ordering cells in CNA heatmaps."""
+
+    EUCLIDEAN = "euclidean"
+    PEARSON = "pearson"
+    SPEARMAN = "spearman"
+
+
+class Genome(StrEnum):
+    """Supported reference genome assemblies."""
+
+    HG20 = "hg20"
+    MM10 = "mm10"
+
+
+class ExecutionBackend(StrEnum):
+    """Supported execution backends."""
+
+    CPU = "cpu"
+    GPU = "gpu"
+    GPU_COMPAT = "gpu-compat"
+
+
+class AnchorStrategy(StrEnum):
+    """How CopyKAT chooses the normal-cell reference."""
+
+    SIGMA = "sigma"
+    MARKERS = "markers"
+
+
+class FinalCallStrategy(StrEnum):
+    """Strategy used to label final copy-number calls."""
+
+    CLUSTERS = "clusters"
+    ARM_CORRELATION = "arm_correlation"
+
+
+class KSMethod(StrEnum):
+    """Breakpoint statistic used during segmentation."""
+
+    MONTE_CARLO = "mc"
+    EXACT = "exact"
+
+
+class GeneIdType(StrEnum):
+    """Gene identifier representation accepted by annotation."""
+
+    SYMBOL = "S"
+    ENSEMBL = "E"
+
+    @classmethod
+    def normalize(cls, value: GeneIdType | str) -> GeneIdType:
+        """Normalize the ID-type argument: E-prefixed values select Ensembl; others select symbols."""
+        if isinstance(value, cls):
+            return value
+        return cls.ENSEMBL if str(value).upper().startswith(cls.ENSEMBL.value) else cls.SYMBOL
+
+
+class CellLineMode(StrEnum):
+    """Whether input contains a pure cell line sample."""
+
+    YES = "yes"
+    NO = "no"
+
+
+class ReportFormat(StrEnum):
+    """Supported output formats for run reports."""
+
+    TEXT = "txt"
+    MARKDOWN = "markdown"
+    HTML = "html"
+    JSON = "json"
+
+
+class ReferenceMode(StrEnum):
+    """Source used to establish the normal reference population."""
+
+    SYNTHETIC = "synthetic"
+    KNOWN_NORMAL = "known_normal"
+    AUTOMATIC = "automatic"
+
+
+class AnchorPath(StrEnum):
+    """Evidence used to choose the normal-cell reference."""
+
+    SIGMA = "sigma"
+    IMMUNE = "immune"
+    ENDOTHELIAL = "endothelial"
+    ENRICHMENT = "enrichment"
+
+
+class PredictionLabel(StrEnum):
+    """Copy-number classification labels emitted by CopyKAT."""
+
+    DIPLOID = "diploid"
+    ANEUPLOID = "aneuploid"
+    DIPLOID_LOW_CONFIDENCE = "c1:diploid:low.conf"
+    ANEUPLOID_LOW_CONFIDENCE = "c2:aneuploid:low.conf"
+    NOT_DEFINED = "not.defined"
+    UNKNOWN = "unknown"
+
+
+class BaselineWarning(StrEnum):
+    """Baseline classification status reported by the pipeline."""
+
+    NONE = ""
+    UNCLASSIFIED = "unclassified.prediction"
+    CELL_LINE = "run with cell line mode"
+    KNOWN_NORMAL = "run with known normal"
+
+
+class DataQualityStatus(StrEnum):
+    """Data-quality status used to choose baseline fallback behavior."""
+
+    OK = "data quality is ok"
+    LOW = "low data quality"
+
+
 class RawInput(TypedDict):
     """Count matrix with its gene and cell names (genes x cells)."""
 
@@ -54,9 +175,6 @@ class RawInput(TypedDict):
 RawMatrix: TypeAlias = pd.DataFrame | npt.NDArray[Any] | SparseMatrix | RawInput | str
 """Inputs accepted by ``copykat()``: genes x cells, or a path to a .mtx/.csv/.tsv file."""
 
-BaselineWarning: TypeAlias = Literal["", "unclassified.prediction", "run with cell line mode", "run with known normal"]
-"""Classification note carried through the pipeline (``WNS`` in the R code)."""
-
 
 class BaselineResult(TypedDict):
     """Reference normal cells and the baseline profile derived from them."""
@@ -66,7 +184,7 @@ class BaselineResult(TypedDict):
     preN: list[str] | IntArray
     """Reference normal cells: names when cell names are known, otherwise column indices."""
     cl: ClusterLabels | None
-    anchor_path: NotRequired[str]
+    anchor_path: NotRequired[AnchorPath]
 
 
 class SyntheticBaselineResult(TypedDict):

copykat_py/allele/orient.py

@@ -13,6 +13,7 @@
 import numpy as np
 import pandas as pd
 
+from copykat_py._types import PredictionLabel
 from copykat_py.allele import counts as _counts
 from copykat_py.allele import hmm as _hmm
 
@@ -21,8 +22,13 @@
 
 
 def _swap_labels(pred):
-    swap = {"aneuploid": "diploid", "diploid": "aneuploid"}
-    return pred.str.replace(r"aneuploid|diploid", lambda m: swap[m.group(0)], regex=True)
+    swap = {
+        PredictionLabel.ANEUPLOID.value: PredictionLabel.DIPLOID.value,
+        PredictionLabel.DIPLOID.value: PredictionLabel.ANEUPLOID.value,
+        PredictionLabel.ANEUPLOID_LOW_CONFIDENCE.value: PredictionLabel.DIPLOID_LOW_CONFIDENCE.value,
+        PredictionLabel.DIPLOID_LOW_CONFIDENCE.value: PredictionLabel.ANEUPLOID_LOW_CONFIDENCE.value,
+    }
+    return pred.replace(swap)
 
 
 def orient_prediction(prediction, counts_dir, phase_csv, min_f_diff=MIN_F_DIFF):
@@ -42,21 +48,27 @@ def orient_prediction(prediction, counts_dir, phase_csv, min_f_diff=MIN_F_DIFF):
     pred = prediction.iloc[:, 1].astype(str)
     idx = prediction.iloc[:, 0].map(_counts.barcode_key).map(col)
     has = idx.notna().to_numpy()
-    groups = {"aneuploid": (pred.str.contains("aneuploid").to_numpy() & has),
-              "diploid": (pred.str.contains("diploid").to_numpy() & ~pred.str.contains("aneuploid").to_numpy() & has)}
-    report = {"snps_phased": int(len(snps)), "flipped": False}
+    is_aneuploid = pred.str.contains(PredictionLabel.ANEUPLOID.value).to_numpy()
+    is_diploid = pred.str.contains(PredictionLabel.DIPLOID.value).to_numpy() & ~is_aneuploid
+    groups = {
+        PredictionLabel.ANEUPLOID: is_aneuploid & has,
+        PredictionLabel.DIPLOID: is_diploid & has,
+    }
+    report = {"snps_phased": len(snps), "flipped": False}
     for g, m in groups.items():
         cols = idx[m].astype(int).to_numpy()
-        report[f"cells_{g}"] = int(len(cols))
+        report[f"cells_{g.value}"] = len(cols)
         if len(cols) < MIN_GROUP_CELLS:
-            report["note"] = f"fewer than {MIN_GROUP_CELLS} {g} cells with allele data; calls left as they are"
+            report["note"] = f"fewer than {MIN_GROUP_CELLS} {g.value} cells with allele data; calls left as they are"
             return prediction.copy(), report
         k = np.asarray(HA[:, cols].sum(axis=1)).ravel()
         n = np.asarray(DP[:, cols].sum(axis=1)).ravel()
         seg, frac = _hmm.segments(snps, k, n)
-        report[f"F_{g}"] = frac
-        report[f"segments_{g}"] = seg.to_dict(orient="records")
-    report["flipped"] = bool(report["F_diploid"] - report["F_aneuploid"] >= min_f_diff)
+        report[f"F_{g.value}"] = frac
+        report[f"segments_{g.value}"] = seg.to_dict(orient="records")
+    report["flipped"] = bool(
+        report[f"F_{PredictionLabel.DIPLOID.value}"] - report[f"F_{PredictionLabel.ANEUPLOID.value}"] >= min_f_diff
+    )
     out = prediction.copy()
     if report["flipped"]:
         out.iloc[:, 1] = _swap_labels(pred).to_numpy()
@@ -68,8 +80,13 @@ def orient_prediction_file(prediction_txt, counts_dir, phase_csv, out_prefix=Non
     (out_prefix defaults to the prediction file's path without 'prediction.txt')."""
     pred = pd.read_csv(prediction_txt, sep="\t")
     out, report = orient_prediction(pred, counts_dir, phase_csv, min_f_diff)
-    prefix = out_prefix if out_prefix is not None else prediction_txt[: -len("prediction.txt")] \
-        if prediction_txt.endswith("prediction.txt") else prediction_txt + "."
+    prefix = (
+        out_prefix
+        if out_prefix is not None
+        else prediction_txt[: -len("prediction.txt")]
+        if prediction_txt.endswith("prediction.txt")
+        else prediction_txt + "."
+    )
     out.to_csv(f"{prefix}prediction.allele_oriented.txt", sep="\t", index=False)
     with open(f"{prefix}allele_orientation.json", "w") as f:
         json.dump(report, f, indent=1)

copykat_py/anchor.py

@@ -26,16 +26,43 @@
 import pandas as pd
 from scipy import sparse
 
+from copykat_py._types import AnchorPath
+
 IMMUNE_MARKERS = ("PTPRC", "LAPTM5", "CORO1A", "CD53", "LCP1", "CD52", "ARHGDIB")
 ENDOTHELIAL_MARKERS = ("PECAM1", "VWF", "CDH5", "CLDN5", "ESAM", "EMCN", "PLVAP")
-MIN_MARKERS = 3          # markers detected for a cell to count as immune/endothelial
-MIN_FRACTION = 0.5       # marked fraction for a cluster to qualify
-MIN_CELLS = 20           # and at least max(MIN_CELLS, 1% of cells)
-ENRICHMENT = 3.0         # fallback: mean immune-marker count vs the rest of the cells
+MIN_MARKERS = 3  # markers detected for a cell to count as immune/endothelial
+MIN_FRACTION = 0.5  # marked fraction for a cluster to qualify
+MIN_CELLS = 20  # and at least max(MIN_CELLS, 1% of cells)
+ENRICHMENT = 3.0  # fallback: mean immune-marker count vs the rest of the cells
 
 # hg38 centromere midpoints (Mb), chr1..22, X (=23)
-_CENTROMERES_MB = np.array([123.4, 93.9, 90.9, 50.0, 48.8, 59.8, 60.1, 45.2, 43.0, 39.8, 53.4, 35.5,
-                            17.7, 17.2, 19.0, 36.8, 25.1, 18.5, 26.2, 28.1, 12.0, 15.0, 60.6])
+_CENTROMERES_MB = np.array(
+    [
+        123.4,
+        93.9,
+        90.9,
+        50.0,
+        48.8,
+        59.8,
+        60.1,
+        45.2,
+        43.0,
+        39.8,
+        53.4,
+        35.5,
+        17.7,
+        17.2,
+        19.0,
+        36.8,
+        25.1,
+        18.5,
+        26.2,
+        28.1,
+        12.0,
+        15.0,
+        60.6,
+    ]
+)
 
 
 def count_markers(rawmat, markers):
@@ -69,14 +96,16 @@ def _largest_population(labels, marked, n_cells):
     return best
 
 
-def choose_anchor_cluster(labels, immune_counts, endothelial_counts, sigma_cluster):
+def choose_anchor_cluster(
+    labels: np.ndarray, immune_counts: np.ndarray, endothelial_counts: np.ndarray, sigma_cluster: int
+) -> tuple[int, AnchorPath]:
     """Return (cluster id, path) for the normal reference; path is "immune", "endothelial",
     "enrichment" or "sigma" and doubles as a coarse confidence indicator."""
     labels = np.asarray(labels)
     imm = np.asarray(immune_counts, dtype=float)
     endo = np.asarray(endothelial_counts, dtype=float)
     n = len(labels)
-    for path, marked in (("immune", imm >= MIN_MARKERS), ("endothelial", endo >= MIN_MARKERS)):
+    for path, marked in ((AnchorPath.IMMUNE, imm >= MIN_MARKERS), (AnchorPath.ENDOTHELIAL, endo >= MIN_MARKERS)):
         k = _largest_population(labels, marked, n)
         if k is not None:
             return k, path
@@ -86,8 +115,8 @@ def choose_anchor_cluster(labels, immune_counts, endothelial_counts, sigma_clust
     m = labels == ids[j]
     rest = imm[~m].mean() if (~m).any() else 0.0
     if m.sum() >= MIN_CELLS and means[j] >= 1.0 and means[j] >= ENRICHMENT * max(rest, 1e-9):
-        return ids[j], "enrichment"
-    return sigma_cluster, "sigma"
+        return int(ids[j]), AnchorPath.ENRICHMENT
+    return sigma_cluster, AnchorPath.SIGMA
 
 
 def _arm_ids(chrom, pos_bp):
@@ -104,10 +133,10 @@ def arm_correlation_calls(values, chrom, pos_bp, anchor_mask, top_fraction=0.10,
     A = np.asarray(agg @ np.asarray(values, dtype=np.float32), dtype=np.float64)  # arms x cells
     A = A.astype(np.float16).astype(np.float64)  # same precision as the evaluated harness
     w = sizes.astype(np.float64)
-    energy = np.sqrt((w[:, None] * A ** 2).sum(0) / w.sum())
+    energy = np.sqrt((w[:, None] * A**2).sum(0) / w.sum())
     consensus = A[:, energy >= np.percentile(energy, 100 * (1 - top_fraction))].mean(1)
     Ac = A - (w[:, None] * A).sum(0) / w.sum()
     cc = consensus - (w * consensus).sum() / w.sum()
-    r = (w[:, None] * Ac * cc[:, None]).sum(0) / np.sqrt((w[:, None] * Ac ** 2).sum(0) * (w * cc ** 2).sum() + 1e-30)
+    r = (w[:, None] * Ac * cc[:, None]).sum(0) / np.sqrt((w[:, None] * Ac**2).sum(0) * (w * cc**2).sum() + 1e-30)
     anchor_mask = np.asarray(anchor_mask, dtype=bool)
     return r > np.percentile(r[anchor_mask], percentile)